diff --git a/jxl/Cargo.toml b/jxl/Cargo.toml index 15a3616..2b906f0 100644 --- a/jxl/Cargo.toml +++ b/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 = [] diff --git a/jxl/src/entropy_coding/huffman.rs b/jxl/src/entropy_coding/huffman.rs index 3cc6444..f86f46d 100644 --- a/jxl/src/entropy_coding/huffman.rs +++ b/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, Error> { + fn decode_simple_table(al_size: usize, br: &mut BitReader) -> Result> { 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, Error> { + ) -> Result> { 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, Error> { + fn build(root_bits: usize, code_lengths: &[u8]) -> Result> { 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 { + pub fn decode(al_size: usize, br: &mut BitReader) -> Result { 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 { + pub fn read(&self, br: &mut BitReader) -> Result { 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,10 +436,10 @@ pub struct HuffmanCodes { } impl HuffmanCodes { - pub fn decode(num: usize, br: &mut BitReader) -> Result { + pub fn decode(num: usize, br: &mut BitReader) -> Result { let alphabet_sizes: Vec = (0..num) .map(|_| Ok(decode_varint16(br)? + 1)) - .collect::>()?; + .collect::>()?; let max = *alphabet_sizes.iter().max().unwrap(); if max as usize > (1 << HUFFMAN_MAX_BITS) { return Err(Error::AlphabetTooLargeHuff(max as usize)); @@ -447,10 +447,10 @@ impl HuffmanCodes { let tables = alphabet_sizes .iter() .map(|sz| Table::decode(*sz as usize, br)) - .collect::>()?; + .collect::>()?; Ok(HuffmanCodes { tables }) } - pub fn read(&self, br: &mut BitReader, ctx: usize) -> Result { + pub fn read(&self, br: &mut BitReader, ctx: usize) -> Result { self.tables[ctx].read(br) } } diff --git a/jxl/src/error.rs b/jxl/src/error.rs index a37392b..2ede5ae 100644 --- a/jxl/src/error.rs +++ b/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 = std::result::Result; diff --git a/jxl/src/image.rs b/jxl/src/image.rs new file mode 100644 index 0000000..0e7a0eb --- /dev/null +++ b/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 { + size: (usize, usize), + data: Vec, +} + +#[derive(Clone, Copy)] +pub struct ImageRect<'a, T: ImageDataType> { + origin: (usize, usize), + size: (usize, usize), + image: &'a Image, +} + +pub struct ImageRectMut<'a, T: ImageDataType> { + origin: (usize, usize), + size: (usize, usize), + image: &'a mut Image, +} + +impl Image { + pub fn new(xsize: usize, ysize: usize) -> Result> { + 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> { + 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> { + 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> { + 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 { + 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::::new(1 << 28, 1 << 28).is_err()); + } + + #[test] + fn rect_basic() -> Result<()> { + let mut image = Image::::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::::new(32, 32)?; + assert!(image.as_rect().to_pgm().starts_with(b"P5\n32 32\n255\n")); + Ok(()) + } +} diff --git a/jxl/src/lib.rs b/jxl/src/lib.rs index c1d77a7..bc36143 100644 --- a/jxl/src/lib.rs +++ b/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;