Add owned and borrowed Image and ImageRect.

jxl/Cargo.toml

@@ -14,3 +14,6 @@ num-traits = "0.2.14"
 array-init = "2.0.0"
 half = "1.7.1"
 jxl_headers_derive = { path = "../jxl_headers_derive" }
+
+[features]
+debug_tools = []

jxl/src/entropy_coding/huffman.rs

@@ -5,7 +5,7 @@
 
 use crate::bit_reader::BitReader;
 use crate::entropy_coding::decode::*;
-use crate::error::Error;
+use crate::error::{Error, Result};
 use crate::util::*;
 
 pub const HUFFMAN_MAX_BITS: usize = 15;
@@ -62,7 +62,7 @@ fn next_table_bit_size(count: &[u16], len: usize, root_bits: usize) -> usize {
 }
 
 impl Table {
-    fn decode_simple_table(al_size: usize, br: &mut BitReader) -> Result<Vec<TableEntry>, Error> {
+    fn decode_simple_table(al_size: usize, br: &mut BitReader) -> Result<Vec<TableEntry>> {
         let max_bits = al_size.ceil_log2();
         let num_symbols = (br.read(2)? + 1) as usize;
         let mut symbols = [0u16; 4];
@@ -198,7 +198,7 @@ impl Table {
         code_length_code_lengths: [u8; CODE_LENGTHS_CODE],
         al_size: usize,
         br: &mut BitReader,
-    ) -> Result<Vec<u8>, Error> {
+    ) -> Result<Vec<u8>> {
         let table = Table::build(5, &code_length_code_lengths)?;
 
         let mut symbol = 0;
@@ -258,7 +258,7 @@ impl Table {
         Ok(code_lengths)
     }
 
-    fn build(root_bits: usize, code_lengths: &[u8]) -> Result<Vec<TableEntry>, Error> {
+    fn build(root_bits: usize, code_lengths: &[u8]) -> Result<Vec<TableEntry>> {
         if code_lengths.len() > 1 << HUFFMAN_MAX_BITS {
             return Err(Error::InvalidHuffman);
         }
@@ -376,7 +376,7 @@ impl Table {
         Ok(table)
     }
 
-    pub fn decode(al_size: usize, br: &mut BitReader) -> Result<Table, Error> {
+    pub fn decode(al_size: usize, br: &mut BitReader) -> Result<Table> {
         let entries = if al_size == 1 {
             vec![TableEntry { bits: 0, value: 0 }; TABLE_SIZE]
         } else {
@@ -416,7 +416,7 @@ impl Table {
         Ok(Table { entries })
     }
 
-    pub fn read(&self, br: &mut BitReader) -> Result<u32, Error> {
+    pub fn read(&self, br: &mut BitReader) -> Result<u32> {
         let mut pos = br.peek(TABLE_BITS) as usize;
         let mut n_bits = self.entries[pos].bits as usize;
         if n_bits > TABLE_BITS {
@@ -436,21 +436,21 @@ pub struct HuffmanCodes {
 }
 
 impl HuffmanCodes {
-    pub fn decode(num: usize, br: &mut BitReader) -> Result<HuffmanCodes, Error> {
+    pub fn decode(num: usize, br: &mut BitReader) -> Result<HuffmanCodes> {
         let alphabet_sizes: Vec<u16> = (0..num)
             .map(|_| Ok(decode_varint16(br)? + 1))
-            .collect::<Result<_, _>>()?;
+            .collect::<Result<_>>()?;
         let max = *alphabet_sizes.iter().max().unwrap();
         if max as usize > (1 << HUFFMAN_MAX_BITS) {
             return Err(Error::AlphabetTooLargeHuff(max as usize));
         }
         let tables = alphabet_sizes
             .iter()
             .map(|sz| Table::decode(*sz as usize, br))
-            .collect::<Result<_, _>>()?;
+            .collect::<Result<_>>()?;
         Ok(HuffmanCodes { tables })
     }
-    pub fn read(&self, br: &mut BitReader, ctx: usize) -> Result<u32, Error> {
+    pub fn read(&self, br: &mut BitReader, ctx: usize) -> Result<u32> {
         self.tables[ctx].read(br)
     }
 }

jxl/src/error.rs

@@ -3,6 +3,8 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+use std::collections::TryReserveError;
+
 use thiserror::Error;
 
 use crate::entropy_coding::huffman::HUFFMAN_MAX_BITS;
@@ -64,4 +66,15 @@ pub enum Error {
     InvalidEcUpsampling(u32, u32, u32),
     #[error("Num_ds: {0} should be smaller than num_passes: {1}")]
     NumPassesTooLarge(u32, u32),
+    #[error("Out of memory: {0}")]
+    OutOfMemory(#[from] TryReserveError),
+    #[error("Image size too large: {0}x{1}")]
+    ImageSizeTooLarge(usize, usize),
+    #[error("Rect out of bounds: {0}x{1}+{2}+{3} rect in {4}x{5} view")]
+    RectOutOfBounds(usize, usize, usize, usize, usize, usize),
+    // Generic arithmetic overflow. Prefer using other errors if possible.
+    #[error("Arithmetic overflow")]
+    ArithmeticOverflow,
 }
+
+pub type Result<T> = std::result::Result<T, Error>;

jxl/src/image.rs

@@ -0,0 +1,257 @@
+use crate::error::{Error, Result};
+
+mod private {
+    pub trait Sealed {}
+}
+
+pub trait ImageDataType: private::Sealed + Copy + Default {
+    fn data_type_id() -> usize;
+}
+
+macro_rules! impl_image_data_type {
+    ($ty: ty, $id: literal) => {
+        impl private::Sealed for $ty {}
+        impl ImageDataType for $ty {
+            fn data_type_id() -> usize {
+                $id
+            }
+        }
+    };
+}
+
+impl_image_data_type!(u8, 0);
+impl_image_data_type!(u16, 1);
+impl_image_data_type!(u32, 2);
+impl_image_data_type!(f32, 3);
+impl_image_data_type!(i8, 4);
+impl_image_data_type!(i16, 5);
+impl_image_data_type!(i32, 6);
+impl_image_data_type!(half::f16, 7);
+
+pub struct Image<T: ImageDataType> {
+    size: (usize, usize),
+    data: Vec<T>,
+}
+
+#[derive(Clone, Copy)]
+pub struct ImageRect<'a, T: ImageDataType> {
+    origin: (usize, usize),
+    size: (usize, usize),
+    image: &'a Image<T>,
+}
+
+pub struct ImageRectMut<'a, T: ImageDataType> {
+    origin: (usize, usize),
+    size: (usize, usize),
+    image: &'a mut Image<T>,
+}
+
+impl<T: ImageDataType> Image<T> {
+    pub fn new(xsize: usize, ysize: usize) -> Result<Image<T>> {
+        let total_size = xsize
+            .checked_mul(ysize)
+            .ok_or(Error::ImageSizeTooLarge(xsize, ysize))?;
+        let mut data = vec![];
+        data.try_reserve_exact(total_size)?;
+        data.resize(total_size, T::default());
+        Ok(Image {
+            size: (xsize, ysize),
+            data,
+        })
+    }
+
+    pub fn size(&self) -> (usize, usize) {
+        self.size
+    }
+
+    pub fn as_rect(&self) -> ImageRect<'_, T> {
+        ImageRect {
+            origin: (0, 0),
+            size: self.size,
+            image: self,
+        }
+    }
+
+    pub fn as_rect_mut(&mut self) -> ImageRectMut<'_, T> {
+        ImageRectMut {
+            origin: (0, 0),
+            size: self.size,
+            image: self,
+        }
+    }
+}
+
+fn rect_size_check(
+    origin: (usize, usize),
+    size: (usize, usize),
+    ssize: (usize, usize),
+) -> Result<()> {
+    if origin
+        .0
+        .checked_add(size.0)
+        .ok_or(Error::ArithmeticOverflow)?
+        > ssize.0
+        || origin
+            .1
+            .checked_add(size.1)
+            .ok_or(Error::ArithmeticOverflow)?
+            > ssize.1
+    {
+        Err(Error::RectOutOfBounds(
+            size.0, size.1, origin.0, origin.1, ssize.0, ssize.1,
+        ))
+    } else {
+        Ok(())
+    }
+}
+
+impl<'a, T: ImageDataType> ImageRect<'a, T> {
+    pub fn rect(
+        &'a self,
+        origin: (usize, usize),
+        size: (usize, usize),
+    ) -> Result<ImageRect<'a, T>> {
+        rect_size_check(origin, size, self.size)?;
+        Ok(ImageRect {
+            origin: (origin.0 + self.origin.0, origin.1 + self.origin.1),
+            size,
+            image: &self.image,
+        })
+    }
+
+    pub fn size(&self) -> (usize, usize) {
+        self.size
+    }
+
+    pub fn row(&self, row: usize) -> &[T] {
+        debug_assert!(row < self.size.1);
+        let start = (row + self.origin.0) * self.image.size.1 + self.origin.1;
+        &self.image.data[start..start + self.size.0]
+    }
+
+    pub fn to_image(&self) -> Result<Image<T>> {
+        let total_size = self.size.0 * self.size.1;
+        let mut data = vec![];
+        data.try_reserve_exact(total_size)?;
+        data.extend((0..self.size.1).flat_map(|x| self.row(x).iter()));
+        Ok(Image {
+            size: self.size,
+            data,
+        })
+    }
+}
+
+impl<'a, T: ImageDataType> ImageRectMut<'a, T> {
+    pub fn rect(
+        &'a mut self,
+        origin: (usize, usize),
+        size: (usize, usize),
+    ) -> Result<ImageRectMut<'a, T>> {
+        rect_size_check(origin, size, self.size)?;
+        Ok(ImageRectMut {
+            origin: (origin.0 + self.origin.0, origin.1 + self.origin.1),
+            size,
+            image: &mut self.image,
+        })
+    }
+
+    pub fn size(&self) -> (usize, usize) {
+        self.size
+    }
+
+    pub fn row(&mut self, row: usize) -> &mut [T] {
+        debug_assert!(row < self.size.1);
+        let start = (row + self.origin.0) * self.image.size.1 + self.origin.1;
+        &mut self.image.data[start..start + self.size.0]
+    }
+
+    pub fn as_rect(&'a self) -> ImageRect<'a, T> {
+        ImageRect {
+            origin: self.origin,
+            size: self.size,
+            image: &self.image,
+        }
+    }
+}
+
+#[cfg(feature = "debug_tools")]
+pub mod debug_tools {
+    use super::{ImageDataType, ImageRect};
+    pub trait ToU8ForWriting {
+        fn to_u8_for_writing(self) -> u8;
+    }
+
+    impl ToU8ForWriting for u8 {
+        fn to_u8_for_writing(self) -> u8 {
+            self
+        }
+    }
+
+    impl ToU8ForWriting for u16 {
+        fn to_u8_for_writing(self) -> u8 {
+            (self as u32 * 255 / 65535) as u8
+        }
+    }
+
+    impl ToU8ForWriting for f32 {
+        fn to_u8_for_writing(self) -> u8 {
+            (self * 255.0).min(255.0).max(0.0).round() as u8
+        }
+    }
+
+    impl ToU8ForWriting for u32 {
+        fn to_u8_for_writing(self) -> u8 {
+            (self as f32 / (2.0f32.powi(32) - 1.0)).to_u8_for_writing()
+        }
+    }
+
+    impl ToU8ForWriting for half::f16 {
+        fn to_u8_for_writing(self) -> u8 {
+            self.to_f32().to_u8_for_writing()
+        }
+    }
+
+    impl<'a, T: ImageDataType + ToU8ForWriting> ImageRect<'a, T> {
+        pub fn to_pgm(&self) -> Vec<u8> {
+            use std::io::Write;
+            let mut ret = vec![];
+            write!(&mut ret, "P5\n{} {}\n255\n", self.size.0, self.size.1).unwrap();
+            ret.extend(
+                (0..self.size.1)
+                    .flat_map(|x| self.row(x).iter())
+                    .map(|x| x.to_u8_for_writing()),
+            );
+            ret
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::error::Result;
+
+    use super::Image;
+
+    #[test]
+    fn huge_image() {
+        assert!(Image::<u8>::new(1 << 28, 1 << 28).is_err());
+    }
+
+    #[test]
+    fn rect_basic() -> Result<()> {
+        let mut image = Image::<u8>::new(32, 32)?;
+        assert_eq!(image.as_rect_mut().rect((0, 0), (1, 1))?.size(), (1, 1));
+        assert!(image.as_rect_mut().rect((30, 30), (3, 3)).is_err());
+        image.as_rect_mut().rect((30, 30), (1, 1))?.row(0)[0] = 1;
+        assert_eq!(image.as_rect_mut().row(30)[30], 1);
+        Ok(())
+    }
+
+    #[cfg(feature = "debug_tools")]
+    #[test]
+    fn to_pgm() -> Result<()> {
+        let image = Image::<u8>::new(32, 32)?;
+        assert!(image.as_rect().to_pgm().starts_with(b"P5\n32 32\n255\n"));
+        Ok(())
+    }
+}

jxl/src/lib.rs

@@ -9,4 +9,5 @@ pub mod entropy_coding;
 pub mod error;
 pub mod headers;
 pub mod icc;
+pub mod image;
 mod util;