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

@@ -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

crates/jxl-frame/src/data/lf_group.rs

@@ -74,14 +74,16 @@ impl<S: Sample> Bundle<LfGroupParams<'_, '_, '_, S>> for LfGroup<S> {
 
         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)

crates/jxl-frame/src/data/pass_group.rs

@@ -147,6 +147,10 @@ pub fn decode_pass_group_modular<S: Sample>(
     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(

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<CutGrid<'g, V>> {
+        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<CutGrid<'g, V>> {
         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) };
 

crates/jxl-modular/src/image.rs

@@ -207,6 +207,7 @@ impl<S: Sample> ModularImageDestination<S> {
         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<S: Sample> ModularImageDestination<S> {
         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<S: Sample> ModularImageDestination<S> {
                     );
                     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<S: Sample> ModularImageDestination<S> {
                     );
                     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<S: Sample> ModularImageDestination<S> {
             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,

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

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");

crates/jxl-oxide/tests/decode.rs

@@ -148,4 +148,5 @@ test! {
     issue_24,
     issue_26,
     issue_32,
+    squeeze_edge,
 }

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,
 );