Rework the decoder so it's safe, add bytemuck dep
This commit is contained in:
parent
6ce544950d
commit
a9a01c6fbd
4 changed files with 79 additions and 76 deletions
|
@ -15,6 +15,9 @@ exclude = [
|
|||
"assets/*",
|
||||
]
|
||||
|
||||
[dependencies]
|
||||
bytemuck = "1.7"
|
||||
|
||||
[workspace]
|
||||
members = ["qoi-bench"]
|
||||
|
||||
|
|
|
@ -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,26 +22,26 @@ 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 pixels = &mut pixels[..];
|
||||
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::<N>::new(); 256];
|
||||
let mut px = Pixel::<N>::new().with_a(0xff);
|
||||
|
||||
loop {
|
||||
match pixels {
|
||||
[px_out, tail @ ..] => {
|
||||
pixels = tail;
|
||||
[px_out, ptail @ ..] => {
|
||||
pixels = ptail;
|
||||
match data {
|
||||
[b1 @ QOI_OP_INDEX..=QOI_OP_INDEX_END, dtail @ ..] => {
|
||||
px = index[usize::from(*b1)];
|
||||
*px_out = px;
|
||||
*px_out = px.into();
|
||||
data = dtail;
|
||||
continue;
|
||||
}
|
||||
|
@ -55,13 +57,11 @@ where
|
|||
data = dtail;
|
||||
}
|
||||
[b1 @ QOI_OP_RUN..=QOI_OP_RUN_END, dtail @ ..] => {
|
||||
*px_out = px.into();
|
||||
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);
|
||||
phead.fill(px.into());
|
||||
pixels = ptail;
|
||||
}
|
||||
data = dtail;
|
||||
continue;
|
||||
}
|
||||
|
@ -89,7 +89,7 @@ where
|
|||
}
|
||||
}
|
||||
index[usize::from(px.hash_index())] = px;
|
||||
*px_out = px;
|
||||
*px_out = px.into();
|
||||
}
|
||||
_ => {
|
||||
cold();
|
||||
|
@ -97,16 +97,8 @@ where
|
|||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
|
|
|
@ -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> {}
|
||||
|
|
|
@ -3,6 +3,7 @@
|
|||
pub const fn cold() {}
|
||||
|
||||
#[inline(always)]
|
||||
#[allow(unused)]
|
||||
pub const fn likely(b: bool) -> bool {
|
||||
if !b {
|
||||
cold()
|
||||
|
|
Loading…
Reference in a new issue