From b461af81accd4f72cfa87774e2b1fbdd8c106ba1 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Mon, 3 Jan 2022 00:47:13 +0300 Subject: [PATCH] Introduce Channels and finish encoder refactor --- src/decode.rs | 46 +++++----- src/encode.rs | 145 +++++++++++++++++++------------- src/error.rs | 6 +- src/header.rs | 55 ++++++++---- src/lib.rs | 6 +- src/{colorspace.rs => types.rs} | 49 +++++++++++ 6 files changed, 205 insertions(+), 102 deletions(-) rename src/{colorspace.rs => types.rs} (55%) diff --git a/src/decode.rs b/src/decode.rs index 197f2ec..59ca777 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -10,6 +10,7 @@ use crate::consts::{ use crate::error::{Error, Result}; use crate::header::Header; use crate::pixel::{Pixel, SupportedChannels}; +use crate::types::Channels; use crate::utils::{cold, unlikely}; const QOI_OP_INDEX_END: u8 = QOI_OP_INDEX | 0x3f; @@ -126,7 +127,7 @@ pub fn qoi_decode_header(data: impl AsRef<[u8]>) -> Result
{ pub struct QoiDecoder<'a> { data: &'a [u8], header: Header, - channels: u8, + channels: Channels, } impl<'a> QoiDecoder<'a> { @@ -139,13 +140,13 @@ impl<'a> QoiDecoder<'a> { } #[inline] - pub const fn with_channels(mut self, channels: u8) -> Self { + pub const fn with_channels(mut self, channels: Channels) -> Self { self.channels = channels; self } #[inline] - pub const fn channels(&self) -> u8 { + pub const fn channels(&self) -> Channels { self.channels } @@ -162,24 +163,23 @@ impl<'a> QoiDecoder<'a> { #[inline] pub fn decode_to_buf(&mut self, mut buf: impl AsMut<[u8]>) -> Result { let buf = buf.as_mut(); - let size = self.header.n_pixels() * self.channels as usize; + let size = self.header.n_pixels() * self.channels.as_u8() as usize; if unlikely(buf.len() < size) { return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size }); } - let n_read = - qoi_decode_impl_slice_all(self.data, buf, self.channels, self.header.channels)?; + let n_read = qoi_decode_impl_slice_all( + self.data, + buf, + self.channels.as_u8(), + self.header.channels.as_u8(), + )?; self.data = &self.data[n_read..]; // can't panic Ok(size) } #[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 mut out = vec![0; self.header.n_pixels() * self.channels.as_u8() as usize]; self.decode_to_buf(&mut out).map(|_| out) } } @@ -272,7 +272,7 @@ fn qoi_decode_impl_stream_all( pub struct QoiStreamDecoder { reader: R, header: Header, - channels: u8, + channels: Channels, } impl QoiStreamDecoder { @@ -284,13 +284,13 @@ impl QoiStreamDecoder { Ok(Self { reader, header, channels: header.channels }) } - pub fn with_channels(mut self, channels: u8) -> Self { + pub fn with_channels(mut self, channels: Channels) -> Self { self.channels = channels; self } #[inline] - pub fn channels(&self) -> u8 { + pub fn channels(&self) -> Channels { self.channels } @@ -312,22 +312,22 @@ impl QoiStreamDecoder { #[inline] pub fn decode_to_buf(&mut self, mut buf: impl AsMut<[u8]>) -> Result { let buf = buf.as_mut(); - let size = self.header.n_pixels() * self.channels as usize; + let size = self.header.n_pixels() * self.channels.as_u8() as usize; if unlikely(buf.len() < size) { return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size }); } - qoi_decode_impl_stream_all(&mut self.reader, buf, self.channels, self.header.channels)?; + qoi_decode_impl_stream_all( + &mut self.reader, + buf, + self.channels.as_u8(), + self.header.channels.as_u8(), + )?; Ok(size) } #[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 mut out = vec![0; self.header.n_pixels() * self.channels.as_u8() as usize]; let _ = self.decode_to_buf(&mut out)?; Ok(out) } diff --git a/src/encode.rs b/src/encode.rs index f0d4e80..12d3300 100644 --- a/src/encode.rs +++ b/src/encode.rs @@ -1,46 +1,25 @@ -use std::convert::TryInto; - -use crate::colorspace::ColorSpace; -use crate::consts::{ - QOI_HEADER_SIZE, QOI_OP_INDEX, QOI_OP_RUN, QOI_PADDING, QOI_PADDING_SIZE, QOI_PIXELS_MAX, -}; +use crate::consts::{QOI_HEADER_SIZE, QOI_OP_INDEX, 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, ColorSpace}; use crate::utils::{unlikely, BytesMut}; #[allow(clippy::cast_possible_truncation)] -fn qoi_encode_impl( - out: &mut [u8], data: &[u8], width: u32, height: u32, colorspace: ColorSpace, -) -> Result +fn qoi_encode_impl(out: &mut [u8], data: &[u8]) -> Result where Pixel: SupportedChannels, { - let max_len = encode_size_required(width, height, N as u8); - if unlikely(out.len() < max_len) { - return Err(Error::OutputBufferTooSmall { size: out.len(), required: max_len }); - } - - let n_pixels = (width as usize) * (height as usize); - if unlikely(data.is_empty()) { - return Err(Error::EmptyImage { width, height }); - } else if unlikely(n_pixels > QOI_PIXELS_MAX) { - return Err(Error::ImageTooLarge { width, height }); - } else if unlikely(n_pixels * N != data.len()) { - return Err(Error::BadEncodingDataSize { size: data.len(), expected: n_pixels * N }); - } - let out_size = out.len(); let mut buf = BytesMut::new(out); - let header = Header { width, height, channels: N as u8, colorspace }; - buf = buf.write_many(&header.encode()); - let mut index = [Pixel::new(); 256]; let mut px_prev = Pixel::new().with_a(0xff); let mut run = 0_u8; let mut px = Pixel::::new().with_a(0xff); + let n_pixels = data.len() / N; + for (i, chunk) in data.chunks_exact(N).enumerate() { px.read(chunk); if px == px_prev { @@ -72,44 +51,94 @@ where } #[inline] -pub fn qoi_encode_to_buf( - mut out: O, data: D, width: u32, height: u32, channels: u8, colorspace: C, -) -> Result -where - O: AsMut<[u8]>, - D: AsRef<[u8]>, - C: TryInto, - Error: From, -{ - let out = out.as_mut(); - let data = data.as_ref(); - let colorspace = colorspace.try_into()?; +fn qoi_encode_impl_all(out: &mut [u8], data: &[u8], channels: Channels) -> Result { match channels { - 3 => qoi_encode_impl::<3>(out, data, width, height, colorspace), - 4 => qoi_encode_impl::<4>(out, data, width, height, colorspace), - _ => Err(Error::InvalidChannels { channels }), + Channels::Rgb => qoi_encode_impl::<3>(out, data), + Channels::Rgba => qoi_encode_impl::<4>(out, data), } } #[inline] -pub fn qoi_encode_to_vec( - data: D, width: u32, height: u32, channels: u8, colorspace: C, -) -> Result> -where - D: AsRef<[u8]>, - C: TryInto, - Error: From, -{ - let size = encode_size_required(width, height, channels); - let mut out = vec![0; size]; // note: we could save time here but that won't be safe anymore - let size = qoi_encode_to_buf(&mut out, data, width, height, channels, colorspace)?; - out.truncate(size); - Ok(out) +pub fn encoded_size_limit(width: u32, height: u32, channels: impl Into) -> usize { + let (width, height) = (width as usize, height as usize); + let n_pixels = width.saturating_mul(height); + QOI_HEADER_SIZE + + n_pixels.saturating_mul(channels.into() as usize) + + n_pixels + + QOI_PADDING_SIZE } #[inline] -pub fn encode_size_required(width: u32, height: u32, channels: u8) -> usize { - let (width, height) = (width as usize, height as usize); - let n_pixels = width.saturating_mul(height); - QOI_HEADER_SIZE + n_pixels.saturating_mul(usize::from(channels)) + n_pixels + QOI_PADDING_SIZE +pub fn qoi_encode_to_buf( + buf: impl AsMut<[u8]>, data: impl AsRef<[u8]>, width: u32, height: u32, +) -> Result { + QoiEncoder::new(&data, width, height)?.encode_to_buf(buf) +} + +#[inline] +pub fn qoi_encode_to_vec(data: impl AsRef<[u8]>, width: u32, height: u32) -> Result> { + QoiEncoder::new(&data, width, height)?.encode_to_vec() +} + +pub struct QoiEncoder<'a> { + data: &'a [u8], + header: Header, +} + +impl<'a> QoiEncoder<'a> { + #[inline] + pub fn new(data: &'a (impl AsRef<[u8]> + ?Sized), width: u32, height: u32) -> Result { + let data = data.as_ref(); + let mut header = + Header::try_new(width, height, Channels::default(), ColorSpace::default())?; + let size = data.len(); + let n_channels = size / header.n_pixels(); + if header.n_pixels() * n_channels != size { + return Err(Error::InvalidImageLength { size, width, height }); + } + header.channels = Channels::try_from(n_channels.min(0xff) as u8)?; + Ok(Self { data, header }) + } + + #[inline] + pub const fn with_colorspace(mut self, colorspace: ColorSpace) -> Self { + self.header = self.header.with_colorspace(colorspace); + self + } + + #[inline] + pub const fn channels(&self) -> Channels { + self.header.channels + } + + #[inline] + pub const fn header(&self) -> &Header { + &self.header + } + + #[inline] + pub fn encoded_size_limit(&self) -> usize { + self.header.encoded_size_limit() + } + + #[inline] + pub fn encode_to_buf(&self, mut buf: impl AsMut<[u8]>) -> Result { + let buf = buf.as_mut(); + let size_required = self.encoded_size_limit(); + if unlikely(buf.len() < size_required) { + return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size_required }); + } + let (head, tail) = buf.split_at_mut(QOI_HEADER_SIZE); // can't panic + head.copy_from_slice(&self.header.encode()); + let n_written = qoi_encode_impl_all(tail, self.data, self.header.channels)?; + Ok(QOI_HEADER_SIZE + n_written) + } + + #[inline] + pub fn encode_to_vec(&self) -> Result> { + let mut out = vec![0_u8; self.encoded_size_limit()]; + let size = self.encode_to_buf(&mut out)?; + out.truncate(size); + Ok(out) + } } diff --git a/src/error.rs b/src/error.rs index 8fda075..02d12f0 100644 --- a/src/error.rs +++ b/src/error.rs @@ -11,7 +11,7 @@ pub enum Error { InvalidChannels { channels: u8 }, EmptyImage { width: u32, height: u32 }, ImageTooLarge { width: u32, height: u32 }, - BadEncodingDataSize { size: usize, expected: usize }, + InvalidImageLength { size: usize, width: u32, height: u32 }, InputBufferTooSmall { size: usize, required: usize }, OutputBufferTooSmall { size: usize, required: usize }, InvalidMagic { magic: u32 }, @@ -36,8 +36,8 @@ impl Display for Error { let mp = QOI_PIXELS_MAX / 1_000_000; write!(f, "image is too large: {}x{} (max={}Mp)", width, height, mp) } - Self::BadEncodingDataSize { size, expected } => { - write!(f, "bad data size when encoding: {} (expected: {})", size, expected) + Self::InvalidImageLength { size, width, height } => { + write!(f, "invalid image length: {} for {}x{}", size, width, height) } Self::InputBufferTooSmall { size, required } => { write!(f, "input buffer size too small: {} (minimum required: {})", size, required) diff --git a/src/header.rs b/src/header.rs index e59e658..c3d1c2a 100644 --- a/src/header.rs +++ b/src/header.rs @@ -1,29 +1,50 @@ +use std::convert::TryInto; + use bytemuck::cast_slice; -use crate::colorspace::ColorSpace; use crate::consts::{QOI_HEADER_SIZE, QOI_MAGIC, QOI_PIXELS_MAX}; +use crate::encoded_size_limit; use crate::error::{Error, Result}; +use crate::types::{Channels, ColorSpace}; use crate::utils::unlikely; #[derive(Copy, Clone, Debug, PartialEq, Eq)] pub struct Header { pub width: u32, pub height: u32, - pub channels: u8, + pub channels: Channels, pub colorspace: ColorSpace, } impl Default for Header { #[inline] fn default() -> Self { - Self { width: 1, height: 1, channels: 3, colorspace: ColorSpace::default() } + Self { + width: 1, + height: 1, + channels: Channels::default(), + colorspace: ColorSpace::default(), + } } } impl Header { #[inline] - pub const fn new(width: u32, height: u32, channels: u8) -> Self { - Self { width, height, channels, colorspace: ColorSpace::Srgb } + pub const fn try_new( + width: u32, height: u32, channels: Channels, colorspace: ColorSpace, + ) -> Result { + if unlikely(height == 0 || width == 0) { + return Err(Error::EmptyImage { width, height }); + } else if unlikely((width as usize).saturating_mul(height as usize) > QOI_PIXELS_MAX) { + return Err(Error::ImageTooLarge { width, height }); + } + Ok(Self { width, height, channels, colorspace }) + } + + #[inline] + pub const fn with_channels(mut self, channels: Channels) -> Self { + self.channels = channels; + self } #[inline] @@ -43,7 +64,7 @@ impl Header { out[..4].copy_from_slice(&QOI_MAGIC.to_be_bytes()); out[4..8].copy_from_slice(&self.width.to_be_bytes()); out[8..12].copy_from_slice(&self.height.to_be_bytes()); - out[12] = self.channels; + out[12] = self.channels.into(); out[13] = self.colorspace.into(); out } @@ -58,22 +79,26 @@ impl Header { let magic = u32::from_be_bytes(v[0]); let width = u32::from_be_bytes(v[1]); let height = u32::from_be_bytes(v[2]); - let channels = data[12]; - let colorspace = ColorSpace::try_from(data[13])?; + let channels = data[12].try_into()?; + let colorspace = data[13].try_into()?; if unlikely(magic != QOI_MAGIC) { return Err(Error::InvalidMagic { magic }); - } else if unlikely(height == 0 || width == 0) { - return Err(Error::EmptyImage { width, height }); - } else if unlikely((width as usize) * (height as usize) > QOI_PIXELS_MAX) { - return Err(Error::ImageTooLarge { width, height }); - } else if unlikely(channels != 3 && channels != 4) { - return Err(Error::InvalidChannels { channels }); } - Ok(Self { width, height, channels, colorspace }) + Self::try_new(width, height, channels, colorspace) } #[inline] pub const fn n_pixels(&self) -> usize { (self.width as usize).saturating_mul(self.height as usize) } + + #[inline] + pub const fn n_bytes(&self) -> usize { + self.n_pixels() * self.channels.as_u8() as usize + } + + #[inline] + pub fn encoded_size_limit(&self) -> usize { + encoded_size_limit(self.width, self.height, self.channels) + } } diff --git a/src/lib.rs b/src/lib.rs index 0f536a0..dd06310 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,21 +9,21 @@ clippy::cargo_common_metadata )] -mod colorspace; mod decode; mod encode; mod error; mod header; mod pixel; +mod types; mod utils; #[doc(hidden)] pub mod consts; -pub use crate::colorspace::ColorSpace; pub use crate::decode::{ qoi_decode_header, qoi_decode_to_buf, qoi_decode_to_vec, QoiDecoder, QoiStreamDecoder, }; -pub use crate::encode::{encode_size_required, qoi_encode_to_buf, qoi_encode_to_vec}; +pub use crate::encode::{encoded_size_limit, qoi_encode_to_buf, qoi_encode_to_vec, QoiEncoder}; pub use crate::error::{Error, Result}; pub use crate::header::Header; +pub use crate::types::{Channels, ColorSpace}; diff --git a/src/colorspace.rs b/src/types.rs similarity index 55% rename from src/colorspace.rs rename to src/types.rs index 57da106..b5c271c 100644 --- a/src/colorspace.rs +++ b/src/types.rs @@ -51,3 +51,52 @@ impl TryFrom for ColorSpace { } } } + +#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)] +#[repr(u8)] +pub enum Channels { + /// Three 8-bit channels (RGB) + Rgb = 3, + /// Four 8-bit channels (RGBA) + Rgba = 4, +} + +impl Channels { + pub const fn is_rgb(self) -> bool { + matches!(self, Self::Rgb) + } + + pub const fn is_rgba(self) -> bool { + matches!(self, Self::Rgba) + } + + pub const fn as_u8(self) -> u8 { + self as u8 + } +} + +impl Default for Channels { + fn default() -> Self { + Self::Rgb + } +} + +impl From for u8 { + #[inline] + fn from(channels: Channels) -> Self { + channels as Self + } +} + +impl TryFrom for Channels { + type Error = Error; + + #[inline] + fn try_from(channels: u8) -> Result { + if unlikely(channels != 3 && channels != 4) { + Err(Error::InvalidChannels { channels }) + } else { + Ok(if channels == 3 { Self::Rgb } else { Self::Rgba }) + } + } +}