Add generic streaming decoder

src/decode.rs

@@ -1,3 +1,5 @@
+use std::io::{Read, Write};
+
 // TODO: can be removed once https://github.com/rust-lang/rust/issues/74985 is stable
 use bytemuck::{cast_slice, cast_slice_mut, Pod};
 
@@ -217,4 +219,179 @@ impl<'a> QoiDecoder<'a> {
         let mut out = vec![0; self.header.n_pixels() * self.channels as usize];
         self.decode_to_buf(&mut out).map(|_| out)
     }
+
+    #[inline]
+    pub fn decode_to_stream<W: Write>(&mut self, writer: &mut W) -> Result<()> {
+        qoi_decode_impl_stream_all(
+            &mut self.data,
+            writer,
+            self.channels,
+            self.header.channels,
+            self.header.n_pixels(),
+        )
+    }
+}
+
+#[inline]
+fn qoi_decode_impl_stream<R: Read, W: Write, const N: usize, const RGBA: bool>(
+    data: &mut R, out: &mut W, mut n_pixels: usize,
+) -> Result<()>
+where
+    Pixel<N>: SupportedChannels,
+    [u8; N]: Pod,
+{
+    let mut index = [[0_u8; N]; 256];
+    let mut px = [0_u8; N];
+    if N == 4 {
+        px[3] = 0xff;
+    }
+
+    while n_pixels != 0 {
+        n_pixels -= 1;
+        let mut p = [0];
+        data.read_exact(&mut p)?;
+        let [b1] = p;
+        match b1 {
+            QOI_OP_INDEX..=QOI_OP_INDEX_END => {
+                px = index[b1 as usize];
+                out.write_all(&px)?;
+                continue;
+            }
+            QOI_OP_RGB => {
+                let mut p = [0; 3];
+                data.read_exact(&mut p)?;
+                decode!(rgb: p[0], p[1], p[2] => px);
+            }
+            QOI_OP_RGBA if RGBA => {
+                let mut p = [0; 4];
+                data.read_exact(&mut p)?;
+                decode!(rgb: p[0], p[1], p[2] => px);
+                if N == 4 {
+                    px[3] = p[3];
+                }
+            }
+            QOI_OP_RUN..=QOI_OP_RUN_END => {
+                let run = (b1 & 0x3f) as usize;
+                for _ in 0..=run {
+                    out.write_all(&px)?;
+                }
+                n_pixels = n_pixels.saturating_sub(run);
+                continue;
+            }
+            QOI_OP_DIFF..=QOI_OP_DIFF_END => {
+                decode!(diff: b1 => px);
+            }
+            QOI_OP_LUMA..=QOI_OP_LUMA_END => {
+                let mut p = [0];
+                data.read_exact(&mut p)?;
+                let [b2] = p;
+                decode!(luma: b1, b2 => px);
+            }
+            _ => {
+                cold();
+            }
+        }
+
+        index[hash_pixel(px) as usize] = px;
+        out.write_all(&px)?;
+    }
+
+    let mut p = [0_u8; QOI_PADDING_SIZE];
+    data.read_exact(&mut p)?;
+    if unlikely(p != QOI_PADDING) {
+        return Err(Error::InvalidPadding);
+    }
+
+    Ok(())
+}
+
+#[inline]
+fn qoi_decode_impl_stream_all<R: Read, W: Write>(
+    data: &mut R, out: &mut W, channels: u8, src_channels: u8, n_pixels: usize,
+) -> Result<()> {
+    match (channels, src_channels) {
+        (3, 3) => qoi_decode_impl_stream::<_, _, 3, false>(data, out, n_pixels),
+        (3, 4) => qoi_decode_impl_stream::<_, _, 3, true>(data, out, n_pixels),
+        (4, 3) => qoi_decode_impl_stream::<_, _, 4, false>(data, out, n_pixels),
+        (4, 4) => qoi_decode_impl_stream::<_, _, 4, true>(data, out, n_pixels),
+        _ => {
+            cold();
+            Err(Error::InvalidChannels { channels })
+        }
+    }
+}
+
+pub struct QoiStreamDecoder<R> {
+    reader: R,
+    header: Header,
+    channels: u8,
+}
+
+impl<R: Read> QoiStreamDecoder<R> {
+    #[inline]
+    pub fn new(mut reader: R) -> Result<Self> {
+        let mut b = [0; QOI_HEADER_SIZE];
+        reader.read_exact(&mut b)?;
+        let header = Header::decode(b)?;
+        Ok(Self { reader, header, channels: header.channels })
+    }
+
+    pub fn with_channels(mut self, channels: u8) -> Self {
+        self.channels = channels;
+        self
+    }
+
+    #[inline]
+    pub fn channels(&self) -> u8 {
+        self.channels
+    }
+
+    #[inline]
+    pub fn header(&self) -> &Header {
+        &self.header
+    }
+
+    #[inline]
+    pub fn reader(&self) -> &R {
+        &self.reader
+    }
+
+    #[inline]
+    pub fn into_reader(self) -> R {
+        self.reader
+    }
+
+    #[inline]
+    pub fn decode_to_buf(&mut self, mut buf: impl AsMut<[u8]>) -> Result<usize> {
+        let mut buf = buf.as_mut();
+        let size = self.header.n_pixels() * self.channels as usize;
+        if unlikely(buf.len() < size) {
+            return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size });
+        }
+        self.decode_to_stream(&mut buf)
+    }
+
+    #[inline]
+    pub fn decode_to_vec(&mut self) -> Result<Vec<u8>> {
+        if unlikely(self.channels > 4) {
+            // prevent accidental over-allocations
+            cold();
+            return Err(Error::InvalidChannels { channels: self.channels });
+        }
+        let mut out = vec![0; self.header.n_pixels() * self.channels as usize];
+        let _ = self.decode_to_stream(&mut out.as_mut_slice())?;
+        Ok(out)
+    }
+
+    #[inline]
+    pub fn decode_to_stream<W: Write>(&mut self, writer: &mut W) -> Result<usize> {
+        qoi_decode_impl_stream_all(
+            &mut self.reader,
+            writer,
+            self.channels,
+            self.header.channels,
+            self.header.n_pixels(),
+        )?;
+        Ok(self.header.n_pixels() * self.channels as usize)
+    }
 }