encoder: add stride & various input formats support

Signed-off-by: Marc-André Lureau <[email protected]>

Author: elmarco

src/encode.rs

@@ -1,6 +1,5 @@
 #[cfg(any(feature = "std", feature = "alloc"))]
 use alloc::{vec, vec::Vec};
-use core::convert::TryFrom;
 #[cfg(feature = "std")]
 use std::io::Write;
 
@@ -10,13 +9,16 @@ use crate::consts::{QOI_HEADER_SIZE, QOI_OP_INDEX, QOI_OP_RUN, QOI_PADDING, QOI_
 use crate::error::{Error, Result};
 use crate::header::Header;
 use crate::pixel::{Pixel, SupportedChannels};
-use crate::types::{Channels, ColorSpace};
+use crate::types::{Channels, ColorSpace, SourceChannels};
 #[cfg(feature = "std")]
 use crate::utils::GenericWriter;
 use crate::utils::{unlikely, BytesMut, Writer};
 
 #[allow(clippy::cast_possible_truncation, unused_assignments, unused_variables)]
-fn encode_impl<W: Writer, const N: usize>(mut buf: W, data: &[u8]) -> Result<usize>
+fn encode_impl<W: Writer, const N: usize, const R: usize>(
+    mut buf: W, data: &[u8], width: usize, height: usize, stride: usize,
+    read_px: impl Fn(&mut Pixel<N>, &[u8]),
+) -> Result<usize>
 where
     Pixel<N>: SupportedChannels,
     [u8; N]: Pod,
@@ -30,59 +32,57 @@ where
     let mut px = Pixel::<N>::new().with_a(0xff);
     let mut index_allowed = false;
 
-    let n_pixels = data.len() / N;
+    let n_pixels = width * height;
 
-    for (i, chunk) in data.chunks_exact(N).enumerate() {
-        px.read(chunk);
-        if px == px_prev {
-            run += 1;
-            if run == 62 || unlikely(i == n_pixels - 1) {
-                buf = buf.write_one(QOI_OP_RUN | (run - 1))?;
-                run = 0;
-            }
-        } else {
-            if run != 0 {
-                #[cfg(not(feature = "reference"))]
-                {
-                    // credits for the original idea: @zakarumych (had to be fixed though)
-                    buf = buf.write_one(if run == 1 && index_allowed {
-                        QOI_OP_INDEX | hash_prev
-                    } else {
-                        QOI_OP_RUN | (run - 1)
-                    })?;
-                }
-                #[cfg(feature = "reference")]
-                {
+    let mut i = 0;
+    for row in data.chunks(stride).take(height) {
+        let pixel_row = &row[..width * R];
+        for chunk in pixel_row.chunks_exact(R) {
+            read_px(&mut px, chunk);
+            if px == px_prev {
+                run += 1;
+                if run == 62 || unlikely(i == n_pixels - 1) {
                     buf = buf.write_one(QOI_OP_RUN | (run - 1))?;
+                    run = 0;
                 }
-                run = 0;
-            }
-            index_allowed = true;
-            let px_rgba = px.as_rgba(0xff);
-            hash_prev = px_rgba.hash_index();
-            let index_px = &mut index[hash_prev as usize];
-            if *index_px == px_rgba {
-                buf = buf.write_one(QOI_OP_INDEX | hash_prev)?;
             } else {
-                *index_px = px_rgba;
-                buf = px.encode_into(px_prev, buf)?;
+                if run != 0 {
+                    #[cfg(not(feature = "reference"))]
+                    {
+                        // credits for the original idea: @zakarumych (had to be fixed though)
+                        buf = buf.write_one(if run == 1 && index_allowed {
+                            QOI_OP_INDEX | hash_prev
+                        } else {
+                            QOI_OP_RUN | (run - 1)
+                        })?;
+                    }
+                    #[cfg(feature = "reference")]
+                    {
+                        buf = buf.write_one(QOI_OP_RUN | (run - 1))?;
+                    }
+                    run = 0;
+                }
+                index_allowed = true;
+                let px_rgba = px.as_rgba(0xff);
+                hash_prev = px_rgba.hash_index();
+                let index_px = &mut index[hash_prev as usize];
+                if *index_px == px_rgba {
+                    buf = buf.write_one(QOI_OP_INDEX | hash_prev)?;
+                } else {
+                    *index_px = px_rgba;
+                    buf = px.encode_into(px_prev, buf)?;
+                }
+                px_prev = px;
             }
-            px_prev = px;
+            i += 1;
         }
     }
 
+    debug_assert_eq!(i, n_pixels);
     buf = buf.write_many(&QOI_PADDING)?;
     Ok(cap.saturating_sub(buf.capacity()))
 }
 
-#[inline]
-fn encode_impl_all<W: Writer>(out: W, data: &[u8], channels: Channels) -> Result<usize> {
-    match channels {
-        Channels::Rgb => encode_impl::<_, 3>(out, data),
-        Channels::Rgba => encode_impl::<_, 4>(out, data),
-    }
-}
-
 /// The maximum number of bytes the encoded image will take.
 ///
 /// Can be used to pre-allocate the buffer to encode the image into.
@@ -113,30 +113,111 @@ pub fn encode_to_vec(data: impl AsRef<[u8]>, width: u32, height: u32) -> Result<
     Encoder::new(&data, width, height)?.encode_to_vec()
 }
 
+/// A builder for creating an encoder.
+pub struct EncoderBuilder<'a> {
+    data: &'a [u8],
+    width: u32,
+    height: u32,
+    stride: Option<usize>,
+    source_channels: Option<SourceChannels>,
+    colorspace: Option<ColorSpace>,
+}
+
+impl<'a> EncoderBuilder<'a> {
+    /// Creates a new encoder builder from a given array of pixel data and image dimensions.
+    pub fn new(data: &'a (impl AsRef<[u8]> + ?Sized), width: u32, height: u32) -> Self {
+        Self {
+            data: data.as_ref(),
+            width,
+            height,
+            stride: None,
+            source_channels: None,
+            colorspace: None,
+        }
+    }
+
+    /// Set the stride of the pixel data.
+    pub const fn stride(mut self, stride: usize) -> Self {
+        self.stride = Some(stride);
+        self
+    }
+
+    /// Set the input format of the pixel data.
+    pub const fn source_channels(mut self, source_channels: SourceChannels) -> Self {
+        self.source_channels = Some(source_channels);
+        self
+    }
+
+    /// Set the colorspace.
+    pub const fn colorspace(mut self, colorspace: ColorSpace) -> Self {
+        self.colorspace = Some(colorspace);
+        self
+    }
+
+    /// Build the encoder.
+    pub fn build(self) -> Result<Encoder<'a>> {
+        let EncoderBuilder { data, width, height, stride, source_channels, colorspace } = self;
+
+        let size = data.len();
+        let no_stride = stride.is_none();
+        let stride = stride.unwrap_or(
+            size.checked_div(height as usize).ok_or(Error::InvalidImageLength {
+                size,
+                width,
+                height,
+            })?,
+        );
+        let source_channels = source_channels.unwrap_or(if stride == width as usize * 3 {
+            SourceChannels::Rgb
+        } else {
+            SourceChannels::Rgba
+        });
+
+        if no_stride {
+            if size != width as usize * height as usize * source_channels.bytes_per_pixel() {
+                return Err(Error::InvalidImageLength { size, width, height });
+            }
+        } else {
+            if stride < width as usize * source_channels.bytes_per_pixel() {
+                return Err(Error::InvalidImageStride { size, width, height, stride });
+            }
+            if stride * (height - 1) as usize + width as usize * source_channels.bytes_per_pixel()
+                < size
+            {
+                return Err(Error::InvalidImageStride { size, width, height, stride });
+            }
+        }
+
+        let channels = source_channels.image_channels();
+        let colorspace = colorspace.unwrap_or_default();
+
+        Ok(Encoder {
+            data,
+            stride,
+            source_channels,
+            header: Header::try_new(self.width, self.height, channels, colorspace)?,
+        })
+    }
+}
+
 /// Encode QOI images into buffers or into streams.
 pub struct Encoder<'a> {
     data: &'a [u8],
+    stride: usize,
+    source_channels: SourceChannels,
     header: Header,
 }
 
 impl<'a> Encoder<'a> {
     /// Creates a new encoder from a given array of pixel data and image dimensions.
+    /// The data must be in RGB(A) order, without fill borders (no stride).
     ///
     /// The number of channels will be inferred automatically (the valid values
     /// are 3 or 4). The color space will be set to sRGB by default.
     #[inline]
     #[allow(clippy::cast_possible_truncation)]
     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 })
+        EncoderBuilder::new(data, width, height).build()
     }
 
     /// Returns a new encoder with modified color space.
@@ -181,7 +262,7 @@ impl<'a> Encoder<'a> {
         }
         let (head, tail) = buf.split_at_mut(QOI_HEADER_SIZE); // can't panic
         head.copy_from_slice(&self.header.encode());
-        let n_written = encode_impl_all(BytesMut::new(tail), self.data, self.header.channels)?;
+        let n_written = self.encode_impl_all(BytesMut::new(tail))?;
         Ok(QOI_HEADER_SIZE + n_written)
     }
 
@@ -203,8 +284,62 @@ impl<'a> Encoder<'a> {
     #[inline]
     pub fn encode_to_stream<W: Write>(&self, writer: &mut W) -> Result<usize> {
         writer.write_all(&self.header.encode())?;
-        let n_written =
-            encode_impl_all(GenericWriter::new(writer), self.data, self.header.channels)?;
+        let n_written = self.encode_impl_all(GenericWriter::new(writer))?;
         Ok(n_written + QOI_HEADER_SIZE)
     }
+
+    #[inline]
+    fn encode_impl_all<W: Writer>(&self, out: W) -> Result<usize> {
+        let width = self.header.width as usize;
+        let height = self.header.height as usize;
+        let stride = self.stride;
+        match self.source_channels {
+            SourceChannels::Rgb => {
+                encode_impl::<_, 3, 3>(out, self.data, width, height, stride, Pixel::read)
+            }
+            SourceChannels::Bgr => {
+                encode_impl::<_, 3, 3>(out, self.data, width, height, stride, |px, c| {
+                    px.update_rgb(c[2], c[1], c[0]);
+                })
+            }
+            SourceChannels::Rgba => {
+                encode_impl::<_, 4, 4>(out, self.data, width, height, stride, Pixel::read)
+            }
+            SourceChannels::Argb => {
+                encode_impl::<_, 4, 4>(out, self.data, width, height, stride, |px, c| {
+                    px.update_rgba(c[1], c[2], c[3], c[0]);
+                })
+            }
+            SourceChannels::Rgbx => {
+                encode_impl::<_, 3, 4>(out, self.data, width, height, stride, |px, c| {
+                    px.read(&c[..3]);
+                })
+            }
+            SourceChannels::Xrgb => {
+                encode_impl::<_, 3, 4>(out, self.data, width, height, stride, |px, c| {
+                    px.update_rgb(c[1], c[2], c[3]);
+                })
+            }
+            SourceChannels::Bgra => {
+                encode_impl::<_, 4, 4>(out, self.data, width, height, stride, |px, c| {
+                    px.update_rgba(c[2], c[1], c[0], c[3]);
+                })
+            }
+            SourceChannels::Abgr => {
+                encode_impl::<_, 4, 4>(out, self.data, width, height, stride, |px, c| {
+                    px.update_rgba(c[3], c[2], c[1], c[0]);
+                })
+            }
+            SourceChannels::Bgrx => {
+                encode_impl::<_, 3, 4>(out, self.data, width, height, stride, |px, c| {
+                    px.update_rgb(c[2], c[1], c[0]);
+                })
+            }
+            SourceChannels::Xbgr => {
+                encode_impl::<_, 4, 4>(out, self.data, width, height, stride, |px, c| {
+                    px.update_rgb(c[3], c[2], c[1]);
+                })
+            }
+        }
+    }
 }

src/error.rs

@@ -16,6 +16,8 @@ pub enum Error {
     InvalidImageDimensions { width: u32, height: u32 },
     /// Image dimensions are inconsistent with image buffer length
     InvalidImageLength { size: usize, width: u32, height: u32 },
+    /// Image stride is inconsistent with image dimension and buffer length
+    InvalidImageStride { size: usize, width: u32, height: u32, stride: usize },
     /// Output buffer is too small to fit encoded/decoded image
     OutputBufferTooSmall { size: usize, required: usize },
     /// Input buffer ended unexpectedly before decoding was finished
@@ -48,6 +50,9 @@ impl Display for Error {
             Self::InvalidImageLength { size, width, height } => {
                 write!(f, "invalid image length: {size} bytes for {width}x{height}")
             }
+            Self::InvalidImageStride { size, width, height, stride } => {
+                write!(f, "invalid image stride: {stride} for {size} bytes of {width}x{height}")
+            }
             Self::OutputBufferTooSmall { size, required } => {
                 write!(f, "output buffer size too small: {size} (required: {required})")
             }

src/lib.rs

@@ -88,8 +88,8 @@ pub use crate::decode::{decode_header, decode_to_buf, Decoder};
 
 #[cfg(any(feature = "alloc", feature = "std"))]
 pub use crate::encode::encode_to_vec;
-pub use crate::encode::{encode_max_len, encode_to_buf, Encoder};
+pub use crate::encode::{encode_max_len, encode_to_buf, Encoder, EncoderBuilder};
 
 pub use crate::error::{Error, Result};
 pub use crate::header::Header;
-pub use crate::types::{Channels, ColorSpace};
+pub use crate::types::{Channels, ColorSpace, SourceChannels};