Implement efficient encoding to generic streams

src/encode.rs

@@ -1,17 +1,18 @@
+use std::io::Write;
+
 use crate::consts::{QOI_HEADER_SIZE, QOI_OP_INDEX, QOI_OP_RUN, QOI_PADDING, QOI_PADDING_SIZE};
 use crate::error::{Error, Result};
 use crate::header::Header;
 use crate::pixel::{Pixel, SupportedChannels};
 use crate::types::{Channels, ColorSpace};
-use crate::utils::{unlikely, BytesMut};
+use crate::utils::{unlikely, BytesMut, GenericWriter, Writer};
 
 #[allow(clippy::cast_possible_truncation)]
-fn qoi_encode_impl<const N: usize>(out: &mut [u8], data: &[u8]) -> Result<usize>
+fn qoi_encode_impl<W: Writer, const N: usize>(mut buf: W, data: &[u8]) -> Result<usize>
 where
     Pixel<N>: SupportedChannels,
 {
-    let out_size = out.len();
-    let mut buf = BytesMut::new(out);
+    let cap = buf.capacity();
 
     let mut index = [Pixel::new(); 256];
     let mut px_prev = Pixel::new().with_a(0xff);
@@ -25,36 +26,36 @@ where
         if px == px_prev {
             run += 1;
             if run == 62 || unlikely(i == n_pixels - 1) {
-                buf = buf.write_one(QOI_OP_RUN | (run - 1));
+                buf = buf.write_one(QOI_OP_RUN | (run - 1))?;
                 run = 0;
             }
         } else {
             if run != 0 {
-                buf = buf.write_one(QOI_OP_RUN | (run - 1));
+                buf = buf.write_one(QOI_OP_RUN | (run - 1))?;
                 run = 0;
             }
             let index_pos = px.hash_index();
             let index_px = &mut index[index_pos as usize];
             let px_rgba = px.as_rgba(0xff);
             if *index_px == px_rgba {
-                buf = buf.write_one(QOI_OP_INDEX | index_pos);
+                buf = buf.write_one(QOI_OP_INDEX | index_pos)?;
             } else {
                 *index_px = px_rgba;
-                buf = px.encode_into(px_prev, buf);
+                buf = px.encode_into(px_prev, buf)?;
             }
             px_prev = px;
         }
     }
 
-    buf = buf.write_many(&QOI_PADDING);
-    Ok(out_size.saturating_sub(buf.len()))
+    buf = buf.write_many(&QOI_PADDING)?;
+    Ok(cap.saturating_sub(buf.capacity()))
 }
 
 #[inline]
-fn qoi_encode_impl_all(out: &mut [u8], data: &[u8], channels: Channels) -> Result<usize> {
+fn qoi_encode_impl_all<W: Writer>(out: W, data: &[u8], channels: Channels) -> Result<usize> {
     match channels {
-        Channels::Rgb => qoi_encode_impl::<3>(out, data),
-        Channels::Rgba => qoi_encode_impl::<4>(out, data),
+        Channels::Rgb => qoi_encode_impl::<_, 3>(out, data),
+        Channels::Rgba => qoi_encode_impl::<_, 4>(out, data),
     }
 }
 
@@ -130,7 +131,7 @@ impl<'a> QoiEncoder<'a> {
         }
         let (head, tail) = buf.split_at_mut(QOI_HEADER_SIZE); // can't panic
         head.copy_from_slice(&self.header.encode());
-        let n_written = qoi_encode_impl_all(tail, self.data, self.header.channels)?;
+        let n_written = qoi_encode_impl_all(BytesMut::new(tail), self.data, self.header.channels)?;
         Ok(QOI_HEADER_SIZE + n_written)
     }
 
@@ -141,4 +142,12 @@ impl<'a> QoiEncoder<'a> {
         out.truncate(size);
         Ok(out)
     }
+
+    #[inline]
+    pub fn encode_to_stream<W: Write>(&self, writer: &mut W) -> Result<usize> {
+        writer.write_all(&self.header.encode())?;
+        let n_written =
+            qoi_encode_impl_all(GenericWriter::new(writer), self.data, self.header.channels)?;
+        Ok(n_written + QOI_HEADER_SIZE)
+    }
 }

src/pixel.rs

@@ -1,4 +1,5 @@
 use crate::consts::{QOI_OP_DIFF, QOI_OP_LUMA, QOI_OP_RGB, QOI_OP_RGBA};
+use crate::error::Result;
 use crate::utils::Writer;
 
 #[derive(Copy, Clone, PartialEq, Eq)]
@@ -118,7 +119,7 @@ impl<const N: usize> Pixel<N> {
     }
 
     #[inline]
-    pub fn encode_into<W: Writer>(&self, px_prev: Self, buf: W) -> W::Output {
+    pub fn encode_into<W: Writer>(&self, px_prev: Self, buf: W) -> Result<W> {
         if N == 3 || self.a_or(0) == px_prev.a_or(0) {
             let vg = self.g().wrapping_sub(px_prev.g());
             let vg_32 = vg.wrapping_add(32);

src/utils.rs

@@ -1,3 +1,7 @@
+use std::io::Write;
+
+use crate::error::Result;
+
 #[inline(always)]
 #[cold]
 pub const fn cold() {}
@@ -19,24 +23,53 @@ pub const fn unlikely(b: bool) -> bool {
     b
 }
 
-pub trait Writer {
-    type Output;
+pub trait Writer: Sized {
+    fn write_one(self, v: u8) -> Result<Self>;
+    fn write_many(self, v: &[u8]) -> Result<Self>;
+    fn capacity(&self) -> usize;
+}
 
-    fn write_one(self, v: u8) -> Self::Output;
-    fn write_many(self, v: &[u8]) -> Self::Output;
+pub struct GenericWriter<W> {
+    writer: W,
+    n_written: usize,
+}
+
+impl<W: Write> GenericWriter<W> {
+    pub fn new(writer: W) -> Self {
+        Self { writer, n_written: 0 }
+    }
+}
+
+impl<W: Write> Writer for GenericWriter<W> {
+    fn write_one(mut self, v: u8) -> Result<Self> {
+        self.n_written += 1;
+        self.writer.write_all(&[v]).map(|_| self).map_err(Into::into)
+    }
+
+    fn write_many(mut self, v: &[u8]) -> Result<Self> {
+        self.n_written += v.len();
+        self.writer.write_all(v).map(|_| self).map_err(Into::into)
+    }
+
+    fn capacity(&self) -> usize {
+        usize::MAX - self.n_written
+    }
 }
 
 impl<'a> Writer for BytesMut<'a> {
-    type Output = BytesMut<'a>;
-
     #[inline]
-    fn write_one(self, v: u8) -> Self {
-        BytesMut::write_one(self, v)
+    fn write_one(self, v: u8) -> Result<Self> {
+        Ok(BytesMut::write_one(self, v))
     }
 
     #[inline]
-    fn write_many(self, v: &[u8]) -> Self {
-        BytesMut::write_many(self, v)
+    fn write_many(self, v: &[u8]) -> Result<Self> {
+        Ok(BytesMut::write_many(self, v))
+    }
+
+    #[inline]
+    fn capacity(&self) -> usize {
+        self.0.len()
     }
 }
 
@@ -64,9 +97,4 @@ impl<'a> BytesMut<'a> {
         head.copy_from_slice(v);
         Self(tail)
     }
-
-    #[inline]
-    pub const fn len(&self) -> usize {
-        self.0.len()
-    }
 }