From 5ad88aa857f35e9334ad01b9f339edf11480179a Mon Sep 17 00:00:00 2001 From: Wonwoo Choi Date: Tue, 27 Feb 2024 18:20:31 +0900 Subject: [PATCH] jxl-modular: Handle zero-sized residual group produced by Squeeze (#258) * jxl-modular: Handle zero-sized residual group produced by Squeeze * jxl-modular: Make it simpler * jxl-frame: Skip empty Modular subimages * jxl-oxide (test): Add decode test `squeeze_edge` * Add an entry to CHANGELOG.md --- CHANGELOG.md | 3 ++ crates/jxl-frame/src/data/lf_group.rs | 18 +++++---- crates/jxl-frame/src/data/pass_group.rs | 4 ++ crates/jxl-grid/src/simple_grid.rs | 24 ++++++++---- crates/jxl-modular/src/image.rs | 36 ++++++++++++++++-- crates/jxl-modular/src/lib.rs | 16 +++++--- crates/jxl-modular/src/transform.rs | 3 +- crates/jxl-oxide/tests/decode | 2 +- crates/jxl-oxide/tests/decode.rs | 1 + crates/jxl-oxide/tests/fuzz_findings.rs | 1 + .../squeeze_groups_off_by_one.fuzz | Bin 0 -> 352 bytes 11 files changed, 83 insertions(+), 25 deletions(-) create mode 100644 crates/jxl-oxide/tests/fuzz_findings/squeeze_groups_off_by_one.fuzz diff --git a/CHANGELOG.md b/CHANGELOG.md index a177e27f..0e69e475 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- `jxl-modular`: Fix Squeeze bug when image dimension is slightly larger than group boundary (#258). + ## [0.7.0] - 2024-02-24 ### Added diff --git a/crates/jxl-frame/src/data/lf_group.rs b/crates/jxl-frame/src/data/lf_group.rs index c674e7ef..e51e74b9 100644 --- a/crates/jxl-frame/src/data/lf_group.rs +++ b/crates/jxl-frame/src/data/lf_group.rs @@ -74,14 +74,16 @@ impl Bundle> for LfGroup { let mut is_mlf_complete = true; if let Some(image) = mlf_group { - let mut subimage = image.recursive(bitstream, global_ma_config, tracker)?; - let mut subimage = subimage.prepare_subimage()?; - subimage.decode( - bitstream, - 1 + frame_header.num_lf_groups() + lf_group_idx, - allow_partial, - )?; - is_mlf_complete = subimage.finish(pool); + if !image.is_empty() { + let mut subimage = image.recursive(bitstream, global_ma_config, tracker)?; + let mut subimage = subimage.prepare_subimage()?; + subimage.decode( + bitstream, + 1 + frame_header.num_lf_groups() + lf_group_idx, + allow_partial, + )?; + is_mlf_complete = subimage.finish(pool); + } } let hf_meta = (frame_header.encoding == Encoding::VarDct && is_mlf_complete) diff --git a/crates/jxl-frame/src/data/pass_group.rs b/crates/jxl-frame/src/data/pass_group.rs index 95d5df24..f0ff2698 100644 --- a/crates/jxl-frame/src/data/pass_group.rs +++ b/crates/jxl-frame/src/data/pass_group.rs @@ -147,6 +147,10 @@ pub fn decode_pass_group_modular( tracker: Option<&AllocTracker>, pool: &JxlThreadPool, ) -> Result<()> { + if modular.is_empty() { + return Ok(()); + } + let mut modular = modular.recursive(bitstream, global_ma_config, tracker)?; let mut subimage = modular.prepare_subimage()?; subimage.decode( diff --git a/crates/jxl-grid/src/simple_grid.rs b/crates/jxl-grid/src/simple_grid.rs index 1bc0ebad..a7ccb194 100644 --- a/crates/jxl-grid/src/simple_grid.rs +++ b/crates/jxl-grid/src/simple_grid.rs @@ -436,6 +436,18 @@ impl<'g, V: Copy> CutGrid<'g, V> { impl<'g, V: Copy> CutGrid<'g, V> { pub fn into_groups(self, group_width: usize, group_height: usize) -> Vec> { + let num_cols = (self.width + group_width - 1) / group_width; + let num_rows = (self.height + group_height - 1) / group_height; + self.into_groups_with_fixed_count(group_width, group_height, num_cols, num_rows) + } + + pub fn into_groups_with_fixed_count( + self, + group_width: usize, + group_height: usize, + num_cols: usize, + num_rows: usize, + ) -> Vec> { let CutGrid { ptr, split_base, @@ -446,15 +458,13 @@ impl<'g, V: Copy> CutGrid<'g, V> { } = self; let split_base = split_base.unwrap_or(ptr.cast()); - let groups_x = (width + group_width - 1) / group_width; - let groups_y = (height + group_height - 1) / group_height; - let mut groups = Vec::with_capacity(groups_x * groups_y); - for gy in 0..groups_y { - let y = gy * group_height; + let mut groups = Vec::with_capacity(num_cols * num_rows); + for gy in 0..num_rows { + let y = (gy * group_height).min(height); let gh = (height - y).min(group_height); let row_ptr = unsafe { ptr.as_ptr().add(y * stride) }; - for gx in 0..groups_x { - let x = gx * group_width; + for gx in 0..num_cols { + let x = (gx * group_width).min(width); let gw = (width - x).min(group_width); let ptr = unsafe { row_ptr.add(x) }; diff --git a/crates/jxl-modular/src/image.rs b/crates/jxl-modular/src/image.rs index 12af52ee..a6606256 100644 --- a/crates/jxl-modular/src/image.rs +++ b/crates/jxl-modular/src/image.rs @@ -207,6 +207,7 @@ impl ModularImageDestination { let num_passes = *pass_shifts.last_key_value().unwrap().0 as usize + 1; let group_dim = self.group_dim; + let group_dim_shift = group_dim.trailing_zeros(); let bit_depth = self.bit_depth; let subimage = self.prepare_subimage()?; let it = subimage @@ -223,7 +224,15 @@ impl ModularImageDestination { let mut pass_groups = Vec::with_capacity(num_passes); pass_groups.resize_with(num_passes, Vec::new); for (i, (info, grid)) in it { - let ModularChannelInfo { hshift, vshift, .. } = info; + let ModularChannelInfo { + original_width, + original_height, + hshift, + vshift, + .. + } = info; + assert!(hshift >= 0 && vshift >= 0); + let grid = match grid { TransformedGrid::Single(g) => g, TransformedGrid::Merged { leader, .. } => leader, @@ -256,7 +265,13 @@ impl ModularImageDestination { ); return Err(crate::Error::InvalidSqueezeParams); } - let grids = grid.into_groups(group_width as usize, group_height as usize); + + let grids = grid.into_groups_with_fixed_count( + group_width as usize, + group_height as usize, + (original_width + group_dim - 1) as usize >> group_dim_shift, + (original_height + group_dim - 1) as usize >> group_dim_shift, + ); (&mut pass_groups[pass_idx], grids) } else { // hshift >= 3 && vshift >= 3 @@ -271,7 +286,12 @@ impl ModularImageDestination { ); return Err(crate::Error::InvalidSqueezeParams); } - let grids = grid.into_groups(lf_group_width as usize, lf_group_height as usize); + let grids = grid.into_groups_with_fixed_count( + lf_group_width as usize, + lf_group_height as usize, + (original_width + (group_dim << 3) - 1) as usize >> (group_dim_shift + 3), + (original_height + (group_dim << 3) - 1) as usize >> (group_dim_shift + 3), + ); (&mut lf_groups, grids) }; @@ -286,9 +306,15 @@ impl ModularImageDestination { for (subimage, grid) in groups.iter_mut().zip(grids) { let width = grid.width() as u32; let height = grid.height() as u32; + if width == 0 || height == 0 { + continue; + } + subimage.channel_info.push(ModularChannelInfo { width, height, + original_width: width << hshift, + original_height: height << vshift, hshift, vshift, }); @@ -369,6 +395,10 @@ impl<'dest, S: Sample> TransformedModularSubimage<'dest, S> { } impl<'dest, S: Sample> TransformedModularSubimage<'dest, S> { + pub fn is_empty(&self) -> bool { + self.channel_info.is_empty() + } + pub fn recursive( self, bitstream: &mut Bitstream, diff --git a/crates/jxl-modular/src/lib.rs b/crates/jxl-modular/src/lib.rs index 6a59dc1c..7b0c2d7a 100644 --- a/crates/jxl-modular/src/lib.rs +++ b/crates/jxl-modular/src/lib.rs @@ -156,27 +156,33 @@ impl ModularChannels { pub struct ModularChannelInfo { width: u32, height: u32, + original_width: u32, + original_height: u32, hshift: i32, vshift: i32, } impl ModularChannelInfo { - fn new(width: u32, height: u32, shift: ChannelShift) -> Self { - let (width, height) = shift.shift_size((width, height)); + fn new(original_width: u32, original_height: u32, shift: ChannelShift) -> Self { + let (width, height) = shift.shift_size((original_width, original_height)); Self { width, height, + original_width, + original_height, hshift: shift.hshift(), vshift: shift.vshift(), } } - fn new_shifted(width: u32, height: u32, hshift: i32, vshift: i32) -> Self { + fn new_unshiftable(width: u32, height: u32) -> Self { Self { width, height, - hshift, - vshift, + original_width: width, + original_height: height, + hshift: -1, + vshift: -1, } } } diff --git a/crates/jxl-modular/src/transform.rs b/crates/jxl-modular/src/transform.rs index 19a9ccef..78d169cc 100644 --- a/crates/jxl-modular/src/transform.rs +++ b/crates/jxl-modular/src/transform.rs @@ -270,7 +270,7 @@ impl Palette { .drain((begin_c as usize + 1)..(end_c as usize)); channels.info.insert( 0, - ModularChannelInfo::new_shifted(self.nb_colours, self.num_c, -1, -1), + ModularChannelInfo::new_unshiftable(self.nb_colours, self.num_c), ); if let Some(grids) = grids { @@ -505,6 +505,7 @@ impl Squeeze { height: h, hshift, vshift, + .. } = ch; if *w == 0 || *h == 0 { tracing::error!(?ch, "Cannot squeeze zero-sized channel"); diff --git a/crates/jxl-oxide/tests/decode b/crates/jxl-oxide/tests/decode index aea2ed3f..241f54b5 160000 --- a/crates/jxl-oxide/tests/decode +++ b/crates/jxl-oxide/tests/decode @@ -1 +1 @@ -Subproject commit aea2ed3f178cc0e9c28c8e4a42afed6017b3d77f +Subproject commit 241f54b56ec116a25be28833fbcd1f4af213096b diff --git a/crates/jxl-oxide/tests/decode.rs b/crates/jxl-oxide/tests/decode.rs index 1c32d17a..b579bcd8 100644 --- a/crates/jxl-oxide/tests/decode.rs +++ b/crates/jxl-oxide/tests/decode.rs @@ -148,4 +148,5 @@ test! { issue_24, issue_26, issue_32, + squeeze_edge, } diff --git a/crates/jxl-oxide/tests/fuzz_findings.rs b/crates/jxl-oxide/tests/fuzz_findings.rs index 4562ab7a..0a5efa01 100644 --- a/crates/jxl-oxide/tests/fuzz_findings.rs +++ b/crates/jxl-oxide/tests/fuzz_findings.rs @@ -74,4 +74,5 @@ test_by_include!( ec_upsampling, too_many_squeezes, prop_abs_overflow, + squeeze_groups_off_by_one, ); diff --git a/crates/jxl-oxide/tests/fuzz_findings/squeeze_groups_off_by_one.fuzz b/crates/jxl-oxide/tests/fuzz_findings/squeeze_groups_off_by_one.fuzz new file mode 100644 index 0000000000000000000000000000000000000000..90b7a516fb780e1658a8dc12c4f5ccc37ad76f29 GIT binary patch literal 352 zcmZusp$@`841GrySWZnqLJ))r637GyECGuks2SV@j^HpTLNEx>)FKdk2sk`nKmr;J zqO{vcCU{BP_xkR7J+^MKw0k}tlTr+{TRtq%K_BA;(=jr*S`2}u2N>yA8pt^Y$7&~< zG0**z;2;68y0~Bs2S5Ln@yjc{ui&U_9kAuPsv1HJlmKsLodQ^6QWi@|&4g@I-P0OM zq$C-@Hh^pL#u8B^kzhl#&{mxO^htK1o18u0N1#1XkUo+v@_j{SSy6oy;7a|drSQOC Dy7W_m literal 0 HcmV?d00001