From 074fc72b2232be96e27e99ffe7779d6385fc9ea8 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sun, 2 Jan 2022 00:03:47 +0300 Subject: [PATCH] Add generic streaming decoder --- src/decode.rs | 177 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/error.rs | 13 +++- 2 files changed, 189 insertions(+), 1 deletion(-) diff --git a/src/decode.rs b/src/decode.rs index eb5769c..782a637 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -1,3 +1,5 @@ +use std::io::{Read, Write}; + // TODO: can be removed once https://github.com/rust-lang/rust/issues/74985 is stable use bytemuck::{cast_slice, cast_slice_mut, Pod}; @@ -217,4 +219,179 @@ impl<'a> QoiDecoder<'a> { let mut out = vec![0; self.header.n_pixels() * self.channels as usize]; self.decode_to_buf(&mut out).map(|_| out) } + + #[inline] + pub fn decode_to_stream(&mut self, writer: &mut W) -> Result<()> { + qoi_decode_impl_stream_all( + &mut self.data, + writer, + self.channels, + self.header.channels, + self.header.n_pixels(), + ) + } +} + +#[inline] +fn qoi_decode_impl_stream( + data: &mut R, out: &mut W, mut n_pixels: usize, +) -> Result<()> +where + Pixel: SupportedChannels, + [u8; N]: Pod, +{ + let mut index = [[0_u8; N]; 256]; + let mut px = [0_u8; N]; + if N == 4 { + px[3] = 0xff; + } + + while n_pixels != 0 { + n_pixels -= 1; + 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]; + out.write_all(&px)?; + continue; + } + QOI_OP_RGB => { + let mut p = [0; 3]; + data.read_exact(&mut p)?; + decode!(rgb: p[0], p[1], p[2] => px); + } + QOI_OP_RGBA if RGBA => { + let mut p = [0; 4]; + data.read_exact(&mut p)?; + decode!(rgb: p[0], p[1], p[2] => px); + if N == 4 { + px[3] = p[3]; + } + } + QOI_OP_RUN..=QOI_OP_RUN_END => { + let run = (b1 & 0x3f) as usize; + for _ in 0..=run { + out.write_all(&px)?; + } + n_pixels = n_pixels.saturating_sub(run); + continue; + } + QOI_OP_DIFF..=QOI_OP_DIFF_END => { + decode!(diff: b1 => px); + } + QOI_OP_LUMA..=QOI_OP_LUMA_END => { + let mut p = [0]; + data.read_exact(&mut p)?; + let [b2] = p; + decode!(luma: b1, b2 => px); + } + _ => { + cold(); + } + } + + index[hash_pixel(px) as usize] = px; + out.write_all(&px)?; + } + + let mut p = [0_u8; QOI_PADDING_SIZE]; + data.read_exact(&mut p)?; + if unlikely(p != QOI_PADDING) { + return Err(Error::InvalidPadding); + } + + Ok(()) +} + +#[inline] +fn qoi_decode_impl_stream_all( + data: &mut R, out: &mut W, channels: u8, src_channels: u8, n_pixels: usize, +) -> Result<()> { + match (channels, src_channels) { + (3, 3) => qoi_decode_impl_stream::<_, _, 3, false>(data, out, n_pixels), + (3, 4) => qoi_decode_impl_stream::<_, _, 3, true>(data, out, n_pixels), + (4, 3) => qoi_decode_impl_stream::<_, _, 4, false>(data, out, n_pixels), + (4, 4) => qoi_decode_impl_stream::<_, _, 4, true>(data, out, n_pixels), + _ => { + cold(); + Err(Error::InvalidChannels { channels }) + } + } +} + +pub struct QoiStreamDecoder { + reader: R, + header: Header, + channels: u8, +} + +impl QoiStreamDecoder { + #[inline] + pub fn new(mut reader: R) -> Result { + let mut b = [0; QOI_HEADER_SIZE]; + reader.read_exact(&mut b)?; + let header = Header::decode(b)?; + Ok(Self { reader, header, channels: header.channels }) + } + + pub fn with_channels(mut self, channels: u8) -> Self { + self.channels = channels; + self + } + + #[inline] + pub fn channels(&self) -> u8 { + self.channels + } + + #[inline] + pub fn header(&self) -> &Header { + &self.header + } + + #[inline] + pub fn reader(&self) -> &R { + &self.reader + } + + #[inline] + pub fn into_reader(self) -> R { + self.reader + } + + #[inline] + pub fn decode_to_buf(&mut self, mut buf: impl AsMut<[u8]>) -> Result { + let mut buf = buf.as_mut(); + let size = self.header.n_pixels() * self.channels as usize; + if unlikely(buf.len() < size) { + return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size }); + } + self.decode_to_stream(&mut buf) + } + + #[inline] + pub fn decode_to_vec(&mut self) -> Result> { + if unlikely(self.channels > 4) { + // prevent accidental over-allocations + cold(); + return Err(Error::InvalidChannels { channels: self.channels }); + } + let mut out = vec![0; self.header.n_pixels() * self.channels as usize]; + let _ = self.decode_to_stream(&mut out.as_mut_slice())?; + Ok(out) + } + + #[inline] + pub fn decode_to_stream(&mut self, writer: &mut W) -> Result { + qoi_decode_impl_stream_all( + &mut self.reader, + writer, + self.channels, + self.header.channels, + self.header.n_pixels(), + )?; + Ok(self.header.n_pixels() * self.channels as usize) + } } diff --git a/src/error.rs b/src/error.rs index 80ae6d9..8fda075 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,11 +1,12 @@ use std::convert::Infallible; use std::error::Error as StdError; use std::fmt::{self, Display}; +use std::io; use std::result::Result as StdResult; use crate::consts::{QOI_MAGIC, QOI_PIXELS_MAX}; -#[derive(Debug, Copy, Clone, PartialEq, Eq)] +#[derive(Debug)] pub enum Error { InvalidChannels { channels: u8 }, EmptyImage { width: u32, height: u32 }, @@ -17,6 +18,7 @@ pub enum Error { UnexpectedBufferEnd, InvalidColorSpace { colorspace: u8 }, InvalidPadding, + IoError(io::Error), } pub type Result = StdResult; @@ -55,6 +57,9 @@ impl Display for Error { Self::InvalidPadding => { write!(f, "invalid padding (stream end marker)") } + Self::IoError(ref err) => { + write!(f, "i/o error: {}", err) + } } } } @@ -66,3 +71,9 @@ impl From for Error { unreachable!() } } + +impl From for Error { + fn from(err: io::Error) -> Self { + Self::IoError(err) + } +}