qoi/src/decode.rs

397 lines
12 KiB
Rust
Raw Normal View History

2022-01-03 14:14:01 +00:00
#[cfg(any(feature = "std", feature = "alloc"))]
use alloc::{vec, vec::Vec};
#[cfg(feature = "std")]
use std::io::Read;
2022-01-01 21:03:47 +00:00
// 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,
QOI_PADDING, QOI_PADDING_SIZE,
};
use crate::error::{Error, Result};
use crate::header::Header;
use crate::pixel::{Pixel, SupportedChannels};
use crate::types::Channels;
use crate::utils::{cold, unlikely};
2021-12-31 10:37:56 +00:00
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;
#[inline]
fn decode_impl_slice<const N: usize, const RGBA: bool>(data: &[u8], out: &mut [u8]) -> Result<usize>
where
Pixel<N>: SupportedChannels,
[u8; N]: Pod,
{
let mut pixels = cast_slice_mut::<_, [u8; N]>(out);
let data_len = data.len();
let mut data = data;
let mut index = [Pixel::<4>::new(); 256];
let mut px = Pixel::<N>::new().with_a(0xff);
let mut px_rgba: Pixel<4>;
2021-12-31 10:37:56 +00:00
while let [px_out, ptail @ ..] = pixels {
pixels = ptail;
match data {
[b1 @ QOI_OP_INDEX..=QOI_OP_INDEX_END, dtail @ ..] => {
px_rgba = index[*b1 as usize];
px.update(px_rgba);
*px_out = px.into();
2021-12-31 10:37:56 +00:00
data = dtail;
continue;
}
[QOI_OP_RGB, r, g, b, dtail @ ..] => {
px.update_rgb(*r, *g, *b);
2021-12-31 10:37:56 +00:00
data = dtail;
}
[QOI_OP_RGBA, r, g, b, a, dtail @ ..] if RGBA => {
px.update_rgba(*r, *g, *b, *a);
2021-12-31 10:37:56 +00:00
data = dtail;
}
[b1 @ QOI_OP_RUN..=QOI_OP_RUN_END, dtail @ ..] => {
*px_out = px.into();
let run = ((b1 & 0x3f) as usize).min(pixels.len());
2021-12-31 10:37:56 +00:00
let (phead, ptail) = pixels.split_at_mut(run); // can't panic
phead.fill(px.into());
2021-12-31 10:37:56 +00:00
pixels = ptail;
data = dtail;
continue;
}
[b1 @ QOI_OP_DIFF..=QOI_OP_DIFF_END, dtail @ ..] => {
px.update_diff(*b1);
2021-12-31 10:37:56 +00:00
data = dtail;
}
[b1 @ QOI_OP_LUMA..=QOI_OP_LUMA_END, b2, dtail @ ..] => {
px.update_luma(*b1, *b2);
2021-12-31 10:37:56 +00:00
data = dtail;
}
_ => {
cold();
if unlikely(data.len() < QOI_PADDING_SIZE) {
2022-01-03 18:12:06 +00:00
return Err(Error::UnexpectedBufferEnd);
2021-12-31 10:37:56 +00:00
}
2021-11-29 11:30:09 +00:00
}
}
px_rgba = px.as_rgba(0xff);
index[px_rgba.hash_index() as usize] = px_rgba;
*px_out = px.into();
}
if unlikely(data.len() < QOI_PADDING_SIZE) {
return Err(Error::UnexpectedBufferEnd);
2022-01-03 11:30:37 +00:00
} else if unlikely(data[..QOI_PADDING_SIZE] != QOI_PADDING) {
return Err(Error::InvalidPadding);
}
Ok(data_len.saturating_sub(data.len()).saturating_sub(QOI_PADDING_SIZE))
}
#[inline]
fn decode_impl_slice_all(
data: &[u8], out: &mut [u8], channels: u8, src_channels: u8,
) -> Result<usize> {
match (channels, src_channels) {
(3, 3) => decode_impl_slice::<3, false>(data, out),
(3, 4) => decode_impl_slice::<3, true>(data, out),
(4, 3) => decode_impl_slice::<4, false>(data, out),
(4, 4) => decode_impl_slice::<4, true>(data, out),
_ => {
cold();
2022-01-01 21:02:55 +00:00
Err(Error::InvalidChannels { channels })
}
}
}
2022-01-03 23:57:42 +00:00
/// Decode the image into a pre-allocated buffer.
///
/// Note: the resulting number of channels will match the header. In order to change
/// the number of channels, use [`Decoder::with_channels`].
#[inline]
pub fn decode_to_buf(buf: impl AsMut<[u8]>, data: impl AsRef<[u8]>) -> Result<Header> {
let mut decoder = Decoder::new(&data)?;
decoder.decode_to_buf(buf)?;
Ok(*decoder.header())
}
2022-01-03 23:57:42 +00:00
/// Decode the image into a newly allocated vector.
///
/// Note: the resulting number of channels will match the header. In order to change
/// the number of channels, use [`Decoder::with_channels`].
2022-01-03 14:14:01 +00:00
#[cfg(any(feature = "std", feature = "alloc"))]
#[inline]
pub fn decode_to_vec(data: impl AsRef<[u8]>) -> Result<(Header, Vec<u8>)> {
let mut decoder = Decoder::new(&data)?;
let out = decoder.decode_to_vec()?;
Ok((*decoder.header(), out))
}
2021-12-01 16:01:58 +00:00
2022-01-03 23:57:42 +00:00
/// Decode the image header from a slice of bytes.
2021-12-01 16:01:58 +00:00
#[inline]
pub fn decode_header(data: impl AsRef<[u8]>) -> Result<Header> {
Header::decode(data)
2021-12-01 16:01:58 +00:00
}
2022-01-03 14:14:01 +00:00
#[cfg(any(feature = "std"))]
2022-01-01 21:03:47 +00:00
#[inline]
fn decode_impl_stream<R: Read, const N: usize, const RGBA: bool>(
data: &mut R, out: &mut [u8],
2022-01-01 21:03:47 +00:00
) -> Result<()>
where
Pixel<N>: SupportedChannels,
[u8; N]: Pod,
{
let mut pixels = cast_slice_mut::<_, [u8; N]>(out);
let mut index = [Pixel::<N>::new(); 256];
let mut px = Pixel::<N>::new().with_a(0xff);
2022-01-01 21:03:47 +00:00
while let [px_out, ptail @ ..] = pixels {
pixels = ptail;
2022-01-01 21:03:47 +00:00
let mut p = [0];
data.read_exact(&mut p)?;
let [b1] = p;
match b1 {
QOI_OP_INDEX..=QOI_OP_INDEX_END => {
px = index[b1 as usize];
*px_out = px.into();
2022-01-01 21:03:47 +00:00
continue;
}
QOI_OP_RGB => {
let mut p = [0; 3];
data.read_exact(&mut p)?;
px.update_rgb(p[0], p[1], p[2]);
2022-01-01 21:03:47 +00:00
}
QOI_OP_RGBA if RGBA => {
let mut p = [0; 4];
data.read_exact(&mut p)?;
px.update_rgba(p[0], p[1], p[2], p[3]);
2022-01-01 21:03:47 +00:00
}
QOI_OP_RUN..=QOI_OP_RUN_END => {
*px_out = px.into();
let run = ((b1 & 0x3f) as usize).min(pixels.len());
let (phead, ptail) = pixels.split_at_mut(run); // can't panic
phead.fill(px.into());
pixels = ptail;
2022-01-01 21:03:47 +00:00
continue;
}
QOI_OP_DIFF..=QOI_OP_DIFF_END => {
px.update_diff(b1);
2022-01-01 21:03:47 +00:00
}
QOI_OP_LUMA..=QOI_OP_LUMA_END => {
let mut p = [0];
data.read_exact(&mut p)?;
let [b2] = p;
px.update_luma(b1, b2);
2022-01-01 21:03:47 +00:00
}
_ => {
cold();
}
}
index[px.hash_index() as usize] = px;
*px_out = px.into();
2022-01-01 21:03:47 +00:00
}
let mut p = [0_u8; QOI_PADDING_SIZE];
data.read_exact(&mut p)?;
if unlikely(p != QOI_PADDING) {
return Err(Error::InvalidPadding);
}
Ok(())
}
2022-01-03 14:14:01 +00:00
#[cfg(feature = "std")]
2022-01-01 21:03:47 +00:00
#[inline]
fn decode_impl_stream_all<R: Read>(
data: &mut R, out: &mut [u8], channels: u8, src_channels: u8,
2022-01-01 21:03:47 +00:00
) -> Result<()> {
match (channels, src_channels) {
(3, 3) => decode_impl_stream::<_, 3, false>(data, out),
(3, 4) => decode_impl_stream::<_, 3, true>(data, out),
(4, 3) => decode_impl_stream::<_, 4, false>(data, out),
(4, 4) => decode_impl_stream::<_, 4, true>(data, out),
2022-01-01 21:03:47 +00:00
_ => {
cold();
Err(Error::InvalidChannels { channels })
}
}
}
2022-01-03 21:25:05 +00:00
#[doc(hidden)]
pub trait Reader: Sized {
fn decode_header(&mut self) -> Result<Header>;
fn decode_image(&mut self, out: &mut [u8], channels: u8, src_channels: u8) -> Result<()>;
}
2022-10-16 23:31:27 +00:00
pub struct Bytes<'a>(&'a [u8]);
2022-01-03 21:25:05 +00:00
impl<'a> Bytes<'a> {
#[inline]
pub const fn new(buf: &'a [u8]) -> Self {
Self(buf)
}
#[inline]
pub const fn as_slice(&self) -> &[u8] {
self.0
}
}
impl<'a> Reader for Bytes<'a> {
#[inline]
fn decode_header(&mut self) -> Result<Header> {
let header = Header::decode(self.0)?;
self.0 = &self.0[QOI_HEADER_SIZE..]; // can't panic
Ok(header)
}
#[inline]
fn decode_image(&mut self, out: &mut [u8], channels: u8, src_channels: u8) -> Result<()> {
let n_read = decode_impl_slice_all(self.0, out, channels, src_channels)?;
self.0 = &self.0[n_read..];
Ok(())
}
}
2022-01-03 14:14:01 +00:00
#[cfg(feature = "std")]
2022-01-03 21:25:05 +00:00
impl<R: Read> Reader for R {
#[inline]
fn decode_header(&mut self) -> Result<Header> {
let mut b = [0; QOI_HEADER_SIZE];
self.read_exact(&mut b)?;
Header::decode(b)
}
#[inline]
fn decode_image(&mut self, out: &mut [u8], channels: u8, src_channels: u8) -> Result<()> {
decode_impl_stream_all(self, out, channels, src_channels)
}
}
2022-01-03 23:57:42 +00:00
/// Decode QOI images from slices or from streams.
2022-01-03 21:25:05 +00:00
#[derive(Clone)]
pub struct Decoder<R> {
2022-01-01 21:03:47 +00:00
reader: R,
header: Header,
channels: Channels,
2022-01-01 21:03:47 +00:00
}
2022-01-03 21:25:05 +00:00
impl<'a> Decoder<Bytes<'a>> {
2022-01-03 23:57:42 +00:00
/// Creates a new decoder from a slice of bytes.
///
/// The header will be decoded immediately upon construction.
///
/// Note: this provides the most efficient decoding, but requires the source data to
/// be loaded in memory in order to decode it. In order to decode from a generic
/// stream, use [`Decoder::from_stream`] instead.
2022-01-03 21:25:05 +00:00
#[inline]
pub fn new(data: &'a (impl AsRef<[u8]> + ?Sized)) -> Result<Self> {
Self::new_impl(Bytes::new(data.as_ref()))
}
2022-01-03 23:57:42 +00:00
/// Returns the undecoded tail of the input slice of bytes.
2022-01-03 21:25:05 +00:00
#[inline]
pub const fn data(&self) -> &[u8] {
self.reader.as_slice()
}
}
2022-01-03 14:14:01 +00:00
#[cfg(feature = "std")]
2022-01-03 21:25:05 +00:00
impl<R: Read> Decoder<R> {
2022-01-03 23:57:42 +00:00
/// Creates a new decoder from a generic reader that implements [`Read`](std::io::Read).
///
/// The header will be decoded immediately upon construction.
///
/// Note: while it's possible to pass a `&[u8]` slice here since it implements `Read`, it
/// would be more efficient to use a specialized constructor instead: [`Decoder::new`].
2022-01-01 21:03:47 +00:00
#[inline]
2022-01-03 21:25:05 +00:00
pub fn from_stream(reader: R) -> Result<Self> {
Self::new_impl(reader)
}
2022-01-03 23:57:42 +00:00
/// Returns an immutable reference to the underlying reader.
2022-01-03 21:25:05 +00:00
#[inline]
2022-10-16 23:34:11 +00:00
pub const fn reader(&self) -> &R {
2022-01-03 21:25:05 +00:00
&self.reader
}
2022-01-03 23:57:42 +00:00
/// Consumes the decoder and returns the underlying reader back.
2022-01-03 21:25:05 +00:00
#[inline]
2022-10-16 23:34:11 +00:00
#[allow(clippy::missing_const_for_fn)]
2022-01-03 21:25:05 +00:00
pub fn into_reader(self) -> R {
self.reader
}
}
impl<R: Reader> Decoder<R> {
#[inline]
fn new_impl(mut reader: R) -> Result<Self> {
let header = reader.decode_header()?;
2022-01-01 21:03:47 +00:00
Ok(Self { reader, header, channels: header.channels })
}
2022-01-03 23:57:42 +00:00
/// Returns a new decoder with modified number of channels.
///
/// By default, the number of channels in the decoded image will be equal
/// to whatever is specified in the header. However, it is also possible
/// to decode RGB into RGBA (in which case the alpha channel will be set
/// to 255), and vice versa (in which case the alpha channel will be ignored).
#[inline]
2022-10-16 23:34:11 +00:00
pub const fn with_channels(mut self, channels: Channels) -> Self {
2022-01-01 21:03:47 +00:00
self.channels = channels;
self
}
2022-01-03 23:57:42 +00:00
/// Returns the number of channels in the decoded image.
///
/// Note: this may differ from the number of channels specified in the header.
2022-01-01 21:03:47 +00:00
#[inline]
2022-10-16 23:34:11 +00:00
pub const fn channels(&self) -> Channels {
2022-01-01 21:03:47 +00:00
self.channels
}
2022-01-03 23:57:42 +00:00
/// Returns the decoded image header.
2022-01-01 21:03:47 +00:00
#[inline]
2022-10-16 23:34:11 +00:00
pub const fn header(&self) -> &Header {
2022-01-01 21:03:47 +00:00
&self.header
}
2022-01-03 23:57:42 +00:00
/// The number of bytes the decoded image will take.
///
/// Can be used to pre-allocate the buffer to decode the image into.
#[inline]
2022-10-16 23:34:11 +00:00
pub const fn required_buf_len(&self) -> usize {
2022-01-03 23:57:42 +00:00
self.header.n_pixels().saturating_mul(self.channels.as_u8() as usize)
}
/// Decodes the image to a pre-allocated buffer and returns the number of bytes written.
///
/// The minimum size of the buffer can be found via [`Decoder::required_buf_len`].
2022-01-01 21:03:47 +00:00
#[inline]
pub fn decode_to_buf(&mut self, mut buf: impl AsMut<[u8]>) -> Result<usize> {
let buf = buf.as_mut();
2022-01-03 23:57:42 +00:00
let size = self.required_buf_len();
2022-01-01 21:03:47 +00:00
if unlikely(buf.len() < size) {
return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size });
}
2022-01-03 21:25:05 +00:00
self.reader.decode_image(buf, self.channels.as_u8(), self.header.channels.as_u8())?;
Ok(size)
2022-01-01 21:03:47 +00:00
}
2022-01-03 23:57:42 +00:00
/// Decodes the image into a newly allocated vector of bytes and returns it.
2022-01-03 21:25:05 +00:00
#[cfg(any(feature = "std", feature = "alloc"))]
2022-01-01 21:03:47 +00:00
#[inline]
pub fn decode_to_vec(&mut self) -> Result<Vec<u8>> {
let mut out = vec![0; self.header.n_pixels() * self.channels.as_u8() as usize];
let _ = self.decode_to_buf(&mut out)?;
2022-01-01 21:03:47 +00:00
Ok(out)
}
}