Encoding - start the big refactor

src/encode.rs

@@ -2,46 +2,12 @@ use std::convert::TryInto;
 
 use crate::colorspace::ColorSpace;
 use crate::consts::{
-    QOI_HEADER_SIZE, QOI_OP_DIFF, QOI_OP_INDEX, QOI_OP_LUMA, QOI_OP_RGB, QOI_OP_RGBA, QOI_OP_RUN,
+    QOI_HEADER_SIZE, QOI_OP_INDEX, QOI_OP_RUN, QOI_PADDING, QOI_PADDING_SIZE, QOI_PIXELS_MAX,
-    QOI_PADDING, QOI_PADDING_SIZE, QOI_PIXELS_MAX,
 };
 use crate::error::{Error, Result};
 use crate::header::Header;
 use crate::pixel::{Pixel, SupportedChannels};
-use crate::utils::{cold, unlikely};
+use crate::utils::{unlikely, BytesMut};
-
-struct WriteBuf<'a> {
-    buf: &'a mut [u8],
-}
-
-impl<'a> WriteBuf<'a> {
-    pub fn new(buf: &'a mut [u8]) -> Self {
-        Self { buf }
-    }
-
-    #[inline]
-    pub fn write<const N: usize>(self, v: [u8; N]) -> Self {
-        let (head, tail) = self.buf.split_at_mut(N);
-        head.copy_from_slice(&v);
-        Self { buf: tail }
-    }
-
-    #[inline]
-    pub fn push(self, v: u8) -> Self {
-        if let Some((first, tail)) = self.buf.split_first_mut() {
-            *first = v;
-            Self { buf: tail }
-        } else {
-            cold();
-            panic!();
-        }
-    }
-
-    #[inline]
-    pub const fn len(&self) -> usize {
-        self.buf.len()
-    }
-}
 
 #[allow(clippy::cast_possible_truncation)]
 fn qoi_encode_impl<const N: usize>(
@@ -65,10 +31,10 @@ where
     }
 
     let out_size = out.len();
-    let mut buf = WriteBuf::new(out);
+    let mut buf = BytesMut::new(out);
 
     let header = Header { width, height, channels: N as u8, colorspace, ..Header::default() };
-    buf = buf.write(header.encode());
+    buf = buf.write_many(&header.encode());
 
     let mut index = [Pixel::new(); 256];
     let mut px_prev = Pixel::new().with_a(0xff);
@@ -80,57 +46,28 @@ where
         if px == px_prev {
             run += 1;
             if run == 62 || unlikely(i == n_pixels - 1) {
-                buf = buf.push(QOI_OP_RUN | (run - 1));
+                buf = buf.write_one(QOI_OP_RUN | (run - 1));
                 run = 0;
             }
         } else {
             if run != 0 {
-                buf = buf.push(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.push(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);
-                if N == 3 || px.a_or(0) == px_prev.a_or(0) {
-                    let vr = px.r().wrapping_sub(px_prev.r());
-                    let vg = px.g().wrapping_sub(px_prev.g());
-                    let vb = px.b().wrapping_sub(px_prev.b());
-
-                    let vg_r = vr.wrapping_sub(vg);
-                    let vg_b = vb.wrapping_sub(vg);
-
-                    let vg_32 = vg.wrapping_add(32);
-
-                    if vg_32 | 63 == 63 {
-                        let (vr_2, vg_2, vb_2) =
-                            (vr.wrapping_add(2), vg.wrapping_add(2), vb.wrapping_add(2));
-                        if vr_2 | vg_2 | vb_2 | 3 == 3 {
-                            buf = buf.push(QOI_OP_DIFF | vr_2 << 4 | vg_2 << 2 | vb_2);
-                        } else {
-                            let (vg_r_8, vg_b_8) = (vg_r.wrapping_add(8), vg_b.wrapping_add(8));
-                            if vg_r_8 | vg_b_8 | 15 == 15 {
-                                buf = buf.write([QOI_OP_LUMA | vg_32, vg_r_8 << 4 | vg_b_8]);
-                            } else {
-                                buf = buf.write([QOI_OP_RGB, px.r(), px.g(), px.b()]);
-                            }
-                        }
-                    } else {
-                        buf = buf.write([QOI_OP_RGB, px.r(), px.g(), px.b()]);
-                    }
-                } else {
-                    buf = buf.write([QOI_OP_RGBA, px.r(), px.g(), px.b(), px.a_or(0xff)]);
-                }
             }
             px_prev = px;
         }
     }
 
-    buf = buf.write(QOI_PADDING);
+    buf = buf.write_many(&QOI_PADDING);
     Ok(out_size.saturating_sub(buf.len()))
 }
 

src/pixel.rs

@@ -1,3 +1,6 @@
+use crate::consts::{QOI_OP_DIFF, QOI_OP_LUMA, QOI_OP_RGB, QOI_OP_RGBA};
+use crate::utils::Writer;
+
 #[derive(Copy, Clone, PartialEq, Eq)]
 #[repr(transparent)]
 pub struct Pixel<const N: usize>(pub [u8; N]);
@@ -66,28 +69,6 @@ impl<const N: usize> Pixel<N> {
         out
     }
 
-    #[inline]
-    pub const fn from_rgb(px: Pixel<3>, with_a: u8) -> Self {
-        let mut i = 0;
-        let mut out = Self::new();
-        while i < 3 {
-            out.0[i] = px.0[i];
-            i += 1;
-        }
-        out.with_a(with_a)
-    }
-
-    #[inline]
-    pub const fn from_array<const M: usize>(arr: [u8; M]) -> Self {
-        let mut i = 0;
-        let mut out = Self::new();
-        while i < N && i < M {
-            out.0[i] = arr[i];
-            i += 1;
-        }
-        out
-    }
-
     #[inline]
     pub const fn r(self) -> u8 {
         self.0[0]
@@ -135,6 +116,36 @@ impl<const N: usize> Pixel<N> {
         self.0[1] = self.0[1].wrapping_add(g);
         self.0[2] = self.0[2].wrapping_add(b);
     }
+
+    #[inline]
+    pub fn encode_into<W: Writer>(&self, px_prev: Self, buf: W) -> W::Output {
+        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);
+            if vg_32 | 63 == 63 {
+                let vr = self.r().wrapping_sub(px_prev.r());
+                let vb = self.b().wrapping_sub(px_prev.b());
+                let vg_r = vr.wrapping_sub(vg);
+                let vg_b = vb.wrapping_sub(vg);
+                let (vr_2, vg_2, vb_2) =
+                    (vr.wrapping_add(2), vg.wrapping_add(2), vb.wrapping_add(2));
+                if vr_2 | vg_2 | vb_2 | 3 == 3 {
+                    buf.write_one(QOI_OP_DIFF | vr_2 << 4 | vg_2 << 2 | vb_2)
+                } else {
+                    let (vg_r_8, vg_b_8) = (vg_r.wrapping_add(8), vg_b.wrapping_add(8));
+                    if vg_r_8 | vg_b_8 | 15 == 15 {
+                        buf.write_many(&[QOI_OP_LUMA | vg_32, vg_r_8 << 4 | vg_b_8])
+                    } else {
+                        buf.write_many(&[QOI_OP_RGB, self.r(), self.g(), self.b()])
+                    }
+                }
+            } else {
+                buf.write_many(&[QOI_OP_RGB, self.r(), self.g(), self.b()])
+            }
+        } else {
+            buf.write_many(&[QOI_OP_RGBA, self.r(), self.g(), self.b(), self.a_or(0xff)])
+        }
+    }
 }
 
 impl<const N: usize> From<Pixel<N>> for [u8; N] {

src/utils.rs

@@ -18,3 +18,55 @@ pub const fn unlikely(b: bool) -> bool {
     }
     b
 }
+
+pub trait Writer {
+    type Output;
+
+    fn write_one(self, v: u8) -> Self::Output;
+    fn write_many(self, v: &[u8]) -> Self::Output;
+}
+
+impl<'a> Writer for BytesMut<'a> {
+    type Output = BytesMut<'a>;
+
+    #[inline]
+    fn write_one(self, v: u8) -> Self {
+        BytesMut::write_one(self, v)
+    }
+
+    #[inline]
+    fn write_many(self, v: &[u8]) -> Self {
+        BytesMut::write_many(self, v)
+    }
+}
+
+pub struct BytesMut<'a>(&'a mut [u8]);
+
+impl<'a> BytesMut<'a> {
+    pub fn new(buf: &'a mut [u8]) -> Self {
+        Self(buf)
+    }
+
+    #[inline]
+    pub fn write_one(self, v: u8) -> Self {
+        if let Some((first, tail)) = self.0.split_first_mut() {
+            *first = v;
+            Self(tail)
+        } else {
+            cold();
+            panic!();
+        }
+    }
+
+    #[inline]
+    pub fn write_many(self, v: &[u8]) -> Self {
+        let (head, tail) = self.0.split_at_mut(v.len());
+        head.copy_from_slice(v);
+        Self(tail)
+    }
+
+    #[inline]
+    pub const fn len(&self) -> usize {
+        self.0.len()
+    }
+}