From a9a01c6fbdf7dd34228d8b358ad75f0c9f4880ba Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Fri, 31 Dec 2021 01:44:59 +0300 Subject: [PATCH] Rework the decoder so it's safe, add bytemuck dep --- Cargo.toml | 3 ++ src/decode.rs | 144 ++++++++++++++++++++++++-------------------------- src/pixel.rs | 7 +++ src/utils.rs | 1 + 4 files changed, 79 insertions(+), 76 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index d555adc..ad1624d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,9 @@ exclude = [ "assets/*", ] +[dependencies] +bytemuck = "1.7" + [workspace] members = ["qoi-bench"] diff --git a/src/decode.rs b/src/decode.rs index d4f4fa3..fe4bf38 100644 --- a/src/decode.rs +++ b/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(data: &[u8], n_pixels: usize) -> Result> where Pixel: 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::::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 pixels = &mut pixels[..]; - let mut data = &data[QOI_HEADER_SIZE..]; - loop { - match pixels { - [px_out, tail @ ..] => { - pixels = tail; - match data { - [b1 @ QOI_OP_INDEX..=QOI_OP_INDEX_END, dtail @ ..] => { - px = index[usize::from(*b1)]; - *px_out = px; - data = dtail; - continue; - } - [QOI_OP_RGB, r, g, b, dtail @ ..] => { - px = Pixel::from_rgb(Pixel::from_array([*r, *g, *b]), px.a_or(0xff)); - data = dtail; - } - [QOI_OP_RGBA, r, g, b, a, dtail @ ..] => { - if N == 4 { - cold(); - } - px = Pixel::from_array([*r, *g, *b, *a]); - data = dtail; - } - [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; - } - _ => { + let mut out = vec![0; n_pixels * N]; // unnecessary zero-init, but w/e + let mut pixels = cast_slice_mut::<_, [u8; N]>(&mut out); + let mut data = &data[QOI_HEADER_SIZE..]; + + let mut index = [Pixel::::new(); 256]; + let mut px = Pixel::::new().with_a(0xff); + + loop { + match pixels { + [px_out, ptail @ ..] => { + pixels = ptail; + match data { + [b1 @ QOI_OP_INDEX..=QOI_OP_INDEX_END, dtail @ ..] => { + px = index[usize::from(*b1)]; + *px_out = px.into(); + data = dtail; + continue; + } + [QOI_OP_RGB, r, g, b, dtail @ ..] => { + px = Pixel::from_rgb(Pixel::from_array([*r, *g, *b]), px.a_or(0xff)); + data = dtail; + } + [QOI_OP_RGBA, r, g, b, a, dtail @ ..] => { + if N == 4 { 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(); - 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) + Ok(out) } pub trait MaybeChannels { diff --git a/src/pixel.rs b/src/pixel.rs index 1a350b2..aa13476 100644 --- a/src/pixel.rs +++ b/src/pixel.rs @@ -102,6 +102,13 @@ impl Pixel { } } +impl From> for [u8; N] { + #[inline(always)] + fn from(px: Pixel) -> Self { + px.0 + } +} + pub trait SupportedChannels {} impl SupportedChannels for Pixel<3> {} diff --git a/src/utils.rs b/src/utils.rs index 9ea837a..56209b3 100644 --- a/src/utils.rs +++ b/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()