Rework the decoder so it's safe, add bytemuck dep

Cargo.toml

@@ -15,6 +15,9 @@ exclude = [
     "assets/*",
 ]
 
+[dependencies]
+bytemuck = "1.7"
+
 [workspace]
 members = ["qoi-bench"]
 

src/decode.rs

@@ -1,4 +1,5 @@
-use std::mem;
+// TODO: can be removed once https://github.com/rust-lang/rust/issues/74985 is stable
+use bytemuck::{cast_slice_mut, Pod};
 
 use crate::consts::{
     QOI_HEADER_SIZE, QOI_OP_DIFF, QOI_OP_INDEX, QOI_OP_LUMA, QOI_OP_RGB, QOI_OP_RGBA, QOI_OP_RUN,
@@ -7,11 +8,12 @@ use crate::consts::{
 use crate::error::{Error, Result};
 use crate::header::Header;
 use crate::pixel::{Pixel, SupportedChannels};
-use crate::utils::{cold, likely, unlikely};
+use crate::utils::{cold, unlikely};
 
 pub fn qoi_decode_impl<const N: usize>(data: &[u8], n_pixels: usize) -> Result<Vec<u8>>
 where
     Pixel<N>: SupportedChannels,
+    [u8; N]: Pod,
 {
     if unlikely(data.len() < QOI_HEADER_SIZE + QOI_PADDING_SIZE) {
         return Err(Error::InputBufferTooSmall {
@@ -20,93 +22,83 @@ where
         });
     }
 
-    let mut pixels = vec![Pixel::<N>::new(); n_pixels];
-    let mut index = [Pixel::new(); 256];
-    let mut px = Pixel::new().with_a(0xff);
-
     const QOI_OP_INDEX_END: u8 = QOI_OP_INDEX | 0x3f;
     const QOI_OP_RUN_END: u8 = QOI_OP_RUN | 0x3d; // <- note, 0x3d (not 0x3f)
     const QOI_OP_DIFF_END: u8 = QOI_OP_DIFF | 0x3f;
     const QOI_OP_LUMA_END: u8 = QOI_OP_LUMA | 0x3f;
 
-    {
+    let mut out = vec![0; n_pixels * N]; // unnecessary zero-init, but w/e
-        let mut pixels = &mut pixels[..];
+    let mut pixels = cast_slice_mut::<_, [u8; N]>(&mut out);
-        let mut data = &data[QOI_HEADER_SIZE..];
+    let mut data = &data[QOI_HEADER_SIZE..];
-        loop {
+
-            match pixels {
+    let mut index = [Pixel::<N>::new(); 256];
-                [px_out, tail @ ..] => {
+    let mut px = Pixel::<N>::new().with_a(0xff);
-                    pixels = tail;
+
-                    match data {
+    loop {
-                        [b1 @ QOI_OP_INDEX..=QOI_OP_INDEX_END, dtail @ ..] => {
+        match pixels {
-                            px = index[usize::from(*b1)];
+            [px_out, ptail @ ..] => {
-                            *px_out = px;
+                pixels = ptail;
-                            data = dtail;
+                match data {
-                            continue;
+                    [b1 @ QOI_OP_INDEX..=QOI_OP_INDEX_END, dtail @ ..] => {
-                        }
+                        px = index[usize::from(*b1)];
-                        [QOI_OP_RGB, r, g, b, dtail @ ..] => {
+                        *px_out = px.into();
-                            px = Pixel::from_rgb(Pixel::from_array([*r, *g, *b]), px.a_or(0xff));
+                        data = dtail;
-                            data = dtail;
+                        continue;
-                        }
+                    }
-                        [QOI_OP_RGBA, r, g, b, a, dtail @ ..] => {
+                    [QOI_OP_RGB, r, g, b, dtail @ ..] => {
-                            if N == 4 {
+                        px = Pixel::from_rgb(Pixel::from_array([*r, *g, *b]), px.a_or(0xff));
-                                cold();
+                        data = dtail;
-                            }
+                    }
-                            px = Pixel::from_array([*r, *g, *b, *a]);
+                    [QOI_OP_RGBA, r, g, b, a, dtail @ ..] => {
-                            data = dtail;
+                        if N == 4 {
-                        }
-                        [b1 @ QOI_OP_RUN..=QOI_OP_RUN_END, dtail @ ..] => {
-                            let run = usize::from(b1 & 0x3f).min(pixels.len());
-                            *px_out = px;
-                            if likely(run != 0) {
-                                let (phead, ptail) = pixels.split_at_mut(run); // can't panic
-                                phead.fill(px);
-                                pixels = ptail;
-                            }
-                            data = dtail;
-                            continue;
-                        }
-                        [b1 @ QOI_OP_DIFF..=QOI_OP_DIFF_END, dtail @ ..] => {
-                            px.rgb_add(
-                                ((b1 >> 4) & 0x03).wrapping_sub(2),
-                                ((b1 >> 2) & 0x03).wrapping_sub(2),
-                                (b1 & 0x03).wrapping_sub(2),
-                            );
-                            data = dtail;
-                        }
-                        [b1 @ QOI_OP_LUMA..=QOI_OP_LUMA_END, b2, dtail @ ..] => {
-                            let vg = (b1 & 0x3f).wrapping_sub(32);
-                            let vg_8 = vg.wrapping_sub(8);
-                            let vr = vg_8.wrapping_add((b2 >> 4) & 0x0f);
-                            let vb = vg_8.wrapping_add(b2 & 0x0f);
-                            px.rgb_add(vr, vg, vb);
-                            data = dtail;
-                        }
-                        _ => {
                             cold();
-                            if unlikely(data.len() < 8) {
+                        }
-                                return Err(Error::UnexpectedBufferEnd);
+                        px = Pixel::from_array([*r, *g, *b, *a]);
-                            }
+                        data = dtail;
+                    }
+                    [b1 @ QOI_OP_RUN..=QOI_OP_RUN_END, dtail @ ..] => {
+                        *px_out = px.into();
+                        let run = usize::from(b1 & 0x3f).min(pixels.len());
+                        let (phead, ptail) = pixels.split_at_mut(run); // can't panic
+                        phead.fill(px.into());
+                        pixels = ptail;
+                        data = dtail;
+                        continue;
+                    }
+                    [b1 @ QOI_OP_DIFF..=QOI_OP_DIFF_END, dtail @ ..] => {
+                        px.rgb_add(
+                            ((b1 >> 4) & 0x03).wrapping_sub(2),
+                            ((b1 >> 2) & 0x03).wrapping_sub(2),
+                            (b1 & 0x03).wrapping_sub(2),
+                        );
+                        data = dtail;
+                    }
+                    [b1 @ QOI_OP_LUMA..=QOI_OP_LUMA_END, b2, dtail @ ..] => {
+                        let vg = (b1 & 0x3f).wrapping_sub(32);
+                        let vg_8 = vg.wrapping_sub(8);
+                        let vr = vg_8.wrapping_add((b2 >> 4) & 0x0f);
+                        let vb = vg_8.wrapping_add(b2 & 0x0f);
+                        px.rgb_add(vr, vg, vb);
+                        data = dtail;
+                    }
+                    _ => {
+                        cold();
+                        if unlikely(data.len() < 8) {
+                            return Err(Error::UnexpectedBufferEnd);
                         }
                     }
-                    index[usize::from(px.hash_index())] = px;
-                    *px_out = px;
-                }
-                _ => {
-                    cold();
-                    break;
                 }
+                index[usize::from(px.hash_index())] = px;
+                *px_out = px.into();
+            }
+            _ => {
+                cold();
+                break;
             }
         }
     }
 
-    let ptr = pixels.as_mut_ptr();
+    Ok(out)
-    mem::forget(pixels);
-    let bytes = unsafe {
-        // Safety: this is safe because we have previously set all the lengths ourselves
-        Vec::from_raw_parts(ptr.cast(), n_pixels * N, n_pixels * N)
-    };
-
-    Ok(bytes)
 }
 
 pub trait MaybeChannels {

src/pixel.rs

@@ -102,6 +102,13 @@ impl<const N: usize> Pixel<N> {
     }
 }
 
+impl<const N: usize> From<Pixel<N>> for [u8; N] {
+    #[inline(always)]
+    fn from(px: Pixel<N>) -> Self {
+        px.0
+    }
+}
+
 pub trait SupportedChannels {}
 
 impl SupportedChannels for Pixel<3> {}

src/utils.rs

@@ -3,6 +3,7 @@
 pub const fn cold() {}
 
 #[inline(always)]
+#[allow(unused)]
 pub const fn likely(b: bool) -> bool {
     if !b {
         cold()