Introduce Channels and finish encoder refactor

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<Header> {
 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<usize> {
         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 =
+        let n_read = qoi_decode_impl_slice_all(
-            qoi_decode_impl_slice_all(self.data, buf, self.channels, self.header.channels)?;
+            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<Vec<u8>> {
-        if unlikely(self.channels > 4) {
+        let mut out = vec![0; self.header.n_pixels() * self.channels.as_u8() as usize];
-            // 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];
         self.decode_to_buf(&mut out).map(|_| out)
     }
 }
@@ -272,7 +272,7 @@ fn qoi_decode_impl_stream_all<R: Read>(
 pub struct QoiStreamDecoder<R> {
     reader: R,
     header: Header,
-    channels: u8,
+    channels: Channels,
 }
 
 impl<R: Read> QoiStreamDecoder<R> {
@@ -284,13 +284,13 @@ impl<R: Read> QoiStreamDecoder<R> {
         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<R: Read> QoiStreamDecoder<R> {
     #[inline]
     pub fn decode_to_buf(&mut self, mut buf: impl AsMut<[u8]>) -> Result<usize> {
         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<Vec<u8>> {
-        if unlikely(self.channels > 4) {
+        let mut out = vec![0; self.header.n_pixels() * self.channels.as_u8() as usize];
-            // 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_buf(&mut out)?;
         Ok(out)
     }

src/encode.rs

@@ -1,46 +1,25 @@
-use std::convert::TryInto;
+use crate::consts::{QOI_HEADER_SIZE, QOI_OP_INDEX, QOI_OP_RUN, QOI_PADDING, QOI_PADDING_SIZE};
-
-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::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<const N: usize>(
+fn qoi_encode_impl<const N: usize>(out: &mut [u8], data: &[u8]) -> Result<usize>
-    out: &mut [u8], data: &[u8], width: u32, height: u32, colorspace: ColorSpace,
-) -> Result<usize>
 where
     Pixel<N>: 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::<N>::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<O, D, C>(
+fn qoi_encode_impl_all(out: &mut [u8], data: &[u8], channels: Channels) -> Result<usize> {
-    mut out: O, data: D, width: u32, height: u32, channels: u8, colorspace: C,
-) -> Result<usize>
-where
-    O: AsMut<[u8]>,
-    D: AsRef<[u8]>,
-    C: TryInto<ColorSpace>,
-    Error: From<C::Error>,
-{
-    let out = out.as_mut();
-    let data = data.as_ref();
-    let colorspace = colorspace.try_into()?;
     match channels {
-        3 => qoi_encode_impl::<3>(out, data, width, height, colorspace),
+        Channels::Rgb => qoi_encode_impl::<3>(out, data),
-        4 => qoi_encode_impl::<4>(out, data, width, height, colorspace),
+        Channels::Rgba => qoi_encode_impl::<4>(out, data),
-        _ => Err(Error::InvalidChannels { channels }),
     }
 }
 
 #[inline]
-pub fn qoi_encode_to_vec<D, C>(
+pub fn encoded_size_limit(width: u32, height: u32, channels: impl Into<u8>) -> usize {
-    data: D, width: u32, height: u32, channels: u8, colorspace: C,
+    let (width, height) = (width as usize, height as usize);
-) -> Result<Vec<u8>>
+    let n_pixels = width.saturating_mul(height);
-where
+    QOI_HEADER_SIZE
-    D: AsRef<[u8]>,
+        + n_pixels.saturating_mul(channels.into() as usize)
-    C: TryInto<ColorSpace>,
+        + n_pixels
-    Error: From<C::Error>,
+        + QOI_PADDING_SIZE
-{
+}
-    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
+#[inline]
-    let size = qoi_encode_to_buf(&mut out, data, width, height, channels, colorspace)?;
+pub fn qoi_encode_to_buf(
+    buf: impl AsMut<[u8]>, data: impl AsRef<[u8]>, width: u32, height: u32,
+) -> Result<usize> {
+    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<Vec<u8>> {
+    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<Self> {
+        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<usize> {
+        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<Vec<u8>> {
+        let mut out = vec![0_u8; self.encoded_size_limit()];
+        let size = self.encode_to_buf(&mut out)?;
         out.truncate(size);
         Ok(out)
     }
-
-#[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
 }

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 } => {
+            Self::InvalidImageLength { size, width, height } => {
-                write!(f, "bad data size when encoding: {} (expected: {})", size, expected)
+                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)

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 {
+    pub const fn try_new(
-        Self { width, height, channels, colorspace: ColorSpace::Srgb }
+        width: u32, height: u32, channels: Channels, colorspace: ColorSpace,
+    ) -> Result<Self> {
+        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 channels = data[12].try_into()?;
-        let colorspace = ColorSpace::try_from(data[13])?;
+        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)
+    }
 }

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};

src/types.rs

@@ -51,3 +51,52 @@ impl TryFrom<u8> 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<Channels> for u8 {
+    #[inline]
+    fn from(channels: Channels) -> Self {
+        channels as Self
+    }
+}
+
+impl TryFrom<u8> for Channels {
+    type Error = Error;
+
+    #[inline]
+    fn try_from(channels: u8) -> Result<Self> {
+        if unlikely(channels != 3 && channels != 4) {
+            Err(Error::InvalidChannels { channels })
+        } else {
+            Ok(if channels == 3 { Self::Rgb } else { Self::Rgba })
+        }
+    }
+}