Rewrite frame decoding

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

@@ -85,12 +85,6 @@ impl Bundle<(&ImageHeader, &FrameHeader)> for LfGlobal {
     }
 }
 
-impl LfGlobal {
-    pub(crate) fn apply_modular_inverse_transform(&mut self) {
-        self.gmodular.modular.inverse_transform();
-    }
-}
-
 define_bundle! {
     #[derive(Debug)]
     pub struct LfGlobalVarDct error(crate::Error) {
@@ -100,7 +94,7 @@ define_bundle! {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct GlobalModular {
     pub ma_config: Option<MaConfig>,
     pub modular: Modular,

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

@@ -1,135 +1,102 @@
-use jxl_bitstream::{read_bits, Bitstream, Bundle};
+use jxl_bitstream::{Bitstream, Bundle};
+use jxl_grid::CutGrid;
 use jxl_modular::{ChannelShift, Modular};
-use jxl_vardct::{HfCoeff, HfCoeffParams};
+use jxl_vardct::{HfCoeffParams, write_hf_coeff};
 
 use crate::{FrameHeader, Result};
 use super::{
     GlobalModular,
-    LfGlobal,
     LfGlobalVarDct,
     LfGroup,
     HfGlobal,
 };
 
-#[derive(Debug, Clone, Copy)]
-pub struct PassGroupParams<'a> {
-    frame_header: &'a FrameHeader,
-    gmodular: &'a GlobalModular,
-    lf_vardct: Option<&'a LfGlobalVarDct>,
-    lf_group: &'a LfGroup,
-    hf_global: Option<&'a HfGlobal>,
-    pass_idx: u32,
-    group_idx: u32,
-    shift: Option<(i32, i32)>,
-}
-
-impl<'a> PassGroupParams<'a> {
-    pub fn new(
-        frame_header: &'a FrameHeader,
-        lf_global: &'a LfGlobal,
-        lf_group: &'a LfGroup,
-        hf_global: Option<&'a HfGlobal>,
-        pass_idx: u32,
-        group_idx: u32,
-        shift: Option<(i32, i32)>,
-    ) -> Self {
-        Self {
-            frame_header,
-            gmodular: &lf_global.gmodular,
-            lf_vardct: lf_global.vardct.as_ref(),
-            lf_group,
-            hf_global,
-            pass_idx,
-            group_idx,
-            shift,
-        }
-    }
+#[derive(Debug)]
+pub struct PassGroupParams<'frame, 'buf, 'g> {
+    pub frame_header: &'frame FrameHeader,
+    pub lf_group: &'frame LfGroup,
+    pub pass_idx: u32,
+    pub group_idx: u32,
+    pub shift: Option<(i32, i32)>,
+    pub gmodular: &'g mut GlobalModular,
+    pub vardct: Option<PassGroupParamsVardct<'frame, 'buf, 'g>>,
 }
 
 #[derive(Debug)]
-pub struct PassGroup {
-    pub hf_coeff: Option<HfCoeff>,
-    pub modular: Modular,
+pub struct PassGroupParamsVardct<'frame, 'buf, 'g> {
+    pub lf_vardct: &'frame LfGlobalVarDct,
+    pub hf_global: &'frame HfGlobal,
+    pub hf_coeff_output: &'buf mut [CutGrid<'g, f32>; 3],
 }
 
-impl Bundle<PassGroupParams<'_>> for PassGroup {
-    type Error = crate::Error;
-
-    fn parse<R: std::io::Read>(bitstream: &mut Bitstream<R>, params: PassGroupParams<'_>) -> Result<Self> {
-        let PassGroupParams {
-            frame_header,
-            gmodular,
-            lf_vardct,
-            lf_group,
-            hf_global,
-            pass_idx,
-            group_idx,
-            shift,
-        } = params;
+pub fn decode_pass_group<R: std::io::Read>(
+    bitstream: &mut Bitstream<R>,
+    params: PassGroupParams,
+) -> Result<()> {
+    let PassGroupParams {
+        frame_header,
+        lf_group,
+        pass_idx,
+        group_idx,
+        shift,
+        gmodular,
+        vardct,
+    } = params;
 
-        let hf_coeff = lf_vardct
-            .zip(lf_group.hf_meta.as_ref())
-            .zip(hf_global)
-            .map(|((lf_vardct, hf_meta), hf_global)| {
-                let hf_pass = &hf_global.hf_passes[pass_idx as usize];
-                let coeff_shift = frame_header.passes.shift.get(pass_idx as usize)
-                    .copied()
-                    .unwrap_or(0);
+    if let (Some(PassGroupParamsVardct { lf_vardct, hf_global, hf_coeff_output }), Some(hf_meta)) = (vardct, &lf_group.hf_meta) {
+        let hf_pass = &hf_global.hf_passes[pass_idx as usize];
+        let coeff_shift = frame_header.passes.shift.get(pass_idx as usize)
+            .copied()
+            .unwrap_or(0);
 
-                let group_col = group_idx % frame_header.groups_per_row();
-                let group_row = group_idx / frame_header.groups_per_row();
-                let lf_col = (group_col % 8) as usize;
-                let lf_row = (group_row % 8) as usize;
-                let group_dim_blocks = (frame_header.group_dim() / 8) as usize;
+        let group_col = group_idx % frame_header.groups_per_row();
+        let group_row = group_idx / frame_header.groups_per_row();
+        let lf_col = (group_col % 8) as usize;
+        let lf_row = (group_row % 8) as usize;
+        let group_dim_blocks = (frame_header.group_dim() / 8) as usize;
 
-                let block_info = &hf_meta.block_info;
+        let block_info = &hf_meta.block_info;
 
-                let block_left = lf_col * group_dim_blocks;
-                let block_top = lf_row * group_dim_blocks;
-                let block_width = (block_info.width() - block_left).min(group_dim_blocks);
-                let block_height = (block_info.height() - block_top).min(group_dim_blocks);
+        let block_left = lf_col * group_dim_blocks;
+        let block_top = lf_row * group_dim_blocks;
+        let block_width = (block_info.width() - block_left).min(group_dim_blocks);
+        let block_height = (block_info.height() - block_top).min(group_dim_blocks);
 
-                let jpeg_upsampling = frame_header.jpeg_upsampling;
-                let block_info = block_info.subgrid(block_left, block_top, block_width, block_height);
-                let lf_quant: Option<[_; 3]> = lf_group.lf_coeff.as_ref().map(|lf_coeff| {
-                    let lf_quant_channels = lf_coeff.lf_quant.image().channel_data();
-                    std::array::from_fn(|idx| {
-                        let lf_quant = &lf_quant_channels[[1, 0, 2][idx]];
-                        let shift = ChannelShift::from_jpeg_upsampling(jpeg_upsampling, idx);
+        let jpeg_upsampling = frame_header.jpeg_upsampling;
+        let block_info = block_info.subgrid(block_left, block_top, block_width, block_height);
+        let lf_quant: Option<[_; 3]> = lf_group.lf_coeff.as_ref().map(|lf_coeff| {
+            let lf_quant_channels = lf_coeff.lf_quant.image().channel_data();
+            std::array::from_fn(|idx| {
+                let lf_quant = &lf_quant_channels[[1, 0, 2][idx]];
+                let shift = ChannelShift::from_jpeg_upsampling(jpeg_upsampling, idx);
 
-                        let block_left = block_left >> shift.hshift();
-                        let block_top = block_top >> shift.vshift();
-                        let (block_width, block_height) = shift.shift_size((block_width as u32, block_height as u32));
-                        lf_quant.subgrid(block_left, block_top, block_width as usize, block_height as usize)
-                    })
-                });
-
-                let params = HfCoeffParams {
-                    num_hf_presets: hf_global.num_hf_presets,
-                    hf_block_ctx: &lf_vardct.hf_block_ctx,
-                    block_info,
-                    jpeg_upsampling,
-                    lf_quant,
-                    hf_pass,
-                    coeff_shift,
-                };
-                HfCoeff::parse(bitstream, params)
+                let block_left = block_left >> shift.hshift();
+                let block_top = block_top >> shift.vshift();
+                let (block_width, block_height) = shift.shift_size((block_width as u32, block_height as u32));
+                lf_quant.subgrid(block_left, block_top, block_width as usize, block_height as usize)
             })
-            .transpose()?;
+        });
 
-        let modular = if let Some((minshift, maxshift)) = shift {
-            let modular_params = gmodular.modular.make_subimage_params_pass_group(gmodular.ma_config.as_ref(), group_idx, minshift, maxshift);
-            let mut modular = read_bits!(bitstream, Bundle(Modular), modular_params)?;
-            modular.decode_image(bitstream, 1 + 3 * frame_header.num_lf_groups() + 17 + pass_idx * frame_header.num_groups() + group_idx)?;
-            modular.inverse_transform();
-            modular
-        } else {
-            Modular::empty()
+        let params = HfCoeffParams {
+            num_hf_presets: hf_global.num_hf_presets,
+            hf_block_ctx: &lf_vardct.hf_block_ctx,
+            block_info,
+            jpeg_upsampling,
+            lf_quant,
+            hf_pass,
+            coeff_shift,
         };
 
-        Ok(Self {
-            hf_coeff,
-            modular,
-        })
+        write_hf_coeff(bitstream, params, hf_coeff_output)?;
+    }
+
+    if let Some((minshift, maxshift)) = shift {
+        let modular_params = gmodular.modular.make_subimage_params_pass_group(gmodular.ma_config.as_ref(), group_idx, minshift, maxshift);
+        let mut modular = Modular::parse(bitstream, modular_params)?;
+        modular.decode_image(bitstream, 1 + 3 * frame_header.num_lf_groups() + 17 + pass_idx * frame_header.num_groups() + group_idx)?;
+        modular.inverse_transform();
+        gmodular.modular.copy_from_modular(modular);
     }
+
+    Ok(())
 }

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

@@ -15,7 +15,8 @@ pub struct Toc {
     num_lf_groups: usize,
     num_groups: usize,
     groups: Vec<TocGroup>,
-    bitstream_order: Vec<usize>,
+    bitstream_to_original: Vec<usize>,
+    original_to_bitstream: Vec<usize>,
     total_size: u64,
 }
 
@@ -38,7 +39,7 @@ impl std::fmt::Debug for Toc {
                 "bitstream_order",
                 &format_args!(
                     "({})",
-                    if self.bitstream_order.is_empty() { "empty" } else { "non-empty" },
+                    if self.bitstream_to_original.is_empty() { "empty" } else { "non-empty" },
                 ),
             )
             .finish_non_exhaustive()
@@ -97,7 +98,7 @@ impl PartialOrd for TocGroupKind {
 impl Toc {
     /// Returns the offset to the beginning of the data.
     pub fn bookmark(&self) -> Bookmark {
-        let idx = self.bitstream_order.first().copied().unwrap_or(0);
+        let idx = self.bitstream_to_original.first().copied().unwrap_or(0);
         self.groups[idx].offset
     }
 
@@ -106,31 +107,22 @@ impl Toc {
         self.groups.len() <= 1
     }
 
-    pub fn lf_global(&self) -> TocGroup {
-        self.groups[0]
-    }
-
-    pub fn lf_group(&self, idx: u32) -> TocGroup {
-        if self.is_single_entry() {
-            panic!("cannot obtain LfGroup offset of single entry frame");
-        } else if (idx as usize) >= self.num_lf_groups {
-            panic!("index out of range: {} >= {} (num_lf_groups)", idx, self.num_lf_groups);
-        } else {
-            self.groups[idx as usize + 1]
-        }
-    }
-
-    pub fn hf_global(&self) -> TocGroup {
-        self.groups[self.num_lf_groups + 1]
-    }
+    pub fn group_index_bitstream_order(&self, kind: TocGroupKind) -> usize {
+        let original_order = match kind {
+            TocGroupKind::All if self.is_single_entry() => 0,
+            _ if self.is_single_entry() => panic!("Cannot request group type of {:?} for single-group frame", kind),
+            TocGroupKind::All => panic!("Cannot request group type of All for multi-group frame"),
+            TocGroupKind::LfGlobal => 0,
+            TocGroupKind::LfGroup(lf_group_idx) => 1 + lf_group_idx as usize,
+            TocGroupKind::HfGlobal => 1 + self.num_lf_groups,
+            TocGroupKind::GroupPass { pass_idx, group_idx } =>
+                1 + self.num_lf_groups + 1 + pass_idx as usize * self.num_groups + group_idx as usize,
+        };
 
-    pub fn pass_group(&self, pass_idx: u32, group_idx: u32) -> TocGroup {
-        if self.is_single_entry() {
-            panic!("cannot obtain PassGroup offset of single entry frame");
+        if self.original_to_bitstream.is_empty() {
+            original_order
         } else {
-            let mut idx = 1 + self.num_lf_groups + 1;
-            idx += (pass_idx as usize * self.num_groups) + group_idx as usize;
-            self.groups[idx]
+            self.original_to_bitstream[original_order]
         }
     }
 
@@ -140,10 +132,10 @@ impl Toc {
     }
 
     pub fn iter_bitstream_order(&self) -> impl Iterator<Item = TocGroup> + Send {
-        let groups = if self.bitstream_order.is_empty() {
+        let groups = if self.bitstream_to_original.is_empty() {
             self.groups.clone()
         } else {
-            self.bitstream_order.iter().map(|&idx| self.groups[idx]).collect()
+            self.bitstream_to_original.iter().map(|&idx| self.groups[idx]).collect()
         };
         groups.into_iter()
     }
@@ -211,18 +203,18 @@ impl Bundle<&crate::FrameHeader> for Toc {
             out
         };
 
-        let (offsets, sizes, bitstream_order) = if permutated_toc {
-            let mut bitstream_order = vec![0usize; permutation.len()];
+        let (offsets, sizes, bitstream_to_original, original_to_bitstream) = if permutated_toc {
+            let mut bitstream_to_original = vec![0usize; permutation.len()];
             let mut offsets_out = Vec::with_capacity(permutation.len());
             let mut sizes_out = Vec::with_capacity(permutation.len());
-            for (idx, perm) in permutation.into_iter().enumerate() {
+            for (idx, &perm) in permutation.iter().enumerate() {
                 offsets_out.push(offsets[perm]);
                 sizes_out.push(sizes[perm]);
-                bitstream_order[perm] = idx;
+                bitstream_to_original[perm] = idx;
             }
-            (offsets_out, sizes_out, bitstream_order)
+            (offsets_out, sizes_out, bitstream_to_original, permutation)
         } else {
-            (offsets, sizes, Vec::new())
+            (offsets, sizes, Vec::new(), Vec::new())
         };
 
         let groups = sizes
@@ -240,7 +232,8 @@ impl Bundle<&crate::FrameHeader> for Toc {
             num_lf_groups: ctx.num_lf_groups() as usize,
             num_groups: num_groups as usize,
             groups,
-            bitstream_order,
+            bitstream_to_original,
+            original_to_bitstream,
             total_size,
         })
     }