jxl_frame/data/
toc.rs

1use crate::Result;
2use jxl_bitstream::{Bitstream, U};
3use jxl_oxide_common::Bundle;
4
5/// Table of contents of a frame.
6///
7/// Frame data are organized in groups. TOC specified the size and order of each group, and it is
8/// decoded after the frame header.
9pub struct Toc {
10    num_lf_groups: usize,
11    num_groups: usize,
12    groups: Vec<TocGroup>,
13    bitstream_to_original: Vec<usize>,
14    original_to_bitstream: Vec<usize>,
15    total_size: usize,
16}
17
18impl std::fmt::Debug for Toc {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        f.debug_struct("Toc")
21            .field("num_lf_groups", &self.num_lf_groups)
22            .field("num_groups", &self.num_groups)
23            .field("total_size", &self.total_size)
24            .field(
25                "groups",
26                &format_args!(
27                    "({} {})",
28                    self.groups.len(),
29                    if self.groups.len() == 1 {
30                        "entry"
31                    } else {
32                        "entries"
33                    },
34                ),
35            )
36            .field(
37                "bitstream_order",
38                &format_args!(
39                    "({})",
40                    if self.bitstream_to_original.is_empty() {
41                        "empty"
42                    } else {
43                        "non-empty"
44                    },
45                ),
46            )
47            .finish_non_exhaustive()
48    }
49}
50
51/// Information about a group in TOC.
52#[derive(Debug, Copy, Clone)]
53pub struct TocGroup {
54    /// Kind of the group.
55    pub kind: TocGroupKind,
56    /// Offset from the beginning of frame header.
57    pub offset: usize,
58    /// Size of the group.
59    pub size: u32,
60}
61
62#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
63pub enum TocGroupKind {
64    All,
65    LfGlobal,
66    LfGroup(u32),
67    HfGlobal,
68    GroupPass { pass_idx: u32, group_idx: u32 },
69}
70
71impl Ord for TocGroupKind {
72    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
73        match (self, other) {
74            (x, y) if x == y => std::cmp::Ordering::Equal,
75            (Self::All, _) => std::cmp::Ordering::Less,
76            (_, Self::All) => std::cmp::Ordering::Greater,
77            (Self::LfGlobal, _) => std::cmp::Ordering::Less,
78            (_, Self::LfGlobal) => std::cmp::Ordering::Greater,
79            (Self::LfGroup(g_self), Self::LfGroup(g_other)) => g_self.cmp(g_other),
80            (Self::LfGroup(_), _) => std::cmp::Ordering::Less,
81            (_, Self::LfGroup(_)) => std::cmp::Ordering::Greater,
82            (Self::HfGlobal, _) => std::cmp::Ordering::Less,
83            (_, Self::HfGlobal) => std::cmp::Ordering::Greater,
84            (
85                Self::GroupPass {
86                    pass_idx: p_self,
87                    group_idx: g_self,
88                },
89                Self::GroupPass {
90                    pass_idx: p_other,
91                    group_idx: g_other,
92                },
93            ) => p_self.cmp(p_other).then(g_self.cmp(g_other)),
94        }
95    }
96}
97
98impl PartialOrd for TocGroupKind {
99    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
100        Some(self.cmp(other))
101    }
102}
103
104impl Toc {
105    /// Returns the offset to the beginning of the data.
106    pub fn bookmark(&self) -> usize {
107        let idx = self.bitstream_to_original.first().copied().unwrap_or(0);
108        self.groups[idx].offset
109    }
110
111    /// Returns whether the frame has only one group.
112    pub fn is_single_entry(&self) -> bool {
113        self.groups.len() <= 1
114    }
115
116    pub fn group_index_bitstream_order(&self, kind: TocGroupKind) -> usize {
117        let original_order = match kind {
118            TocGroupKind::All if self.is_single_entry() => 0,
119            _ if self.is_single_entry() => {
120                panic!("Cannot request group type of {kind:?} for single-group frame",)
121            }
122            TocGroupKind::All => panic!("Cannot request group type of All for multi-group frame"),
123            TocGroupKind::LfGlobal => 0,
124            TocGroupKind::LfGroup(lf_group_idx) => 1 + lf_group_idx as usize,
125            TocGroupKind::HfGlobal => 1 + self.num_lf_groups,
126            TocGroupKind::GroupPass {
127                pass_idx,
128                group_idx,
129            } => {
130                1 + self.num_lf_groups
131                    + 1
132                    + pass_idx as usize * self.num_groups
133                    + group_idx as usize
134            }
135        };
136
137        if self.original_to_bitstream.is_empty() {
138            original_order
139        } else {
140            self.original_to_bitstream[original_order]
141        }
142    }
143
144    /// Returns the total size of the frame data in bytes.
145    pub fn total_byte_size(&self) -> usize {
146        self.total_size
147    }
148
149    pub fn iter_bitstream_order(&self) -> impl Iterator<Item = TocGroup> + Send {
150        let groups = if self.bitstream_to_original.is_empty() {
151            self.groups.clone()
152        } else {
153            self.bitstream_to_original
154                .iter()
155                .map(|&idx| self.groups[idx])
156                .collect()
157        };
158        groups.into_iter()
159    }
160}
161
162impl Toc {
163    pub(crate) fn adjust_offsets(&mut self, global_frame_offset: usize) {
164        if global_frame_offset == 0 {
165            return;
166        }
167
168        for group in &mut self.groups {
169            group.offset = group
170                .offset
171                .checked_sub(global_frame_offset)
172                .expect("group offset is smaller than global frame offset");
173        }
174    }
175}
176
177impl Bundle<&crate::FrameHeader> for Toc {
178    type Error = crate::Error;
179
180    fn parse(bitstream: &mut Bitstream, ctx: &crate::FrameHeader) -> Result<Self> {
181        let num_groups = ctx.num_groups();
182        let num_passes = ctx.passes.num_passes;
183
184        let entry_count = if num_groups == 1 && num_passes == 1 {
185            1
186        } else {
187            1 + ctx.num_lf_groups() + 1 + num_groups * num_passes
188        };
189
190        if entry_count > 65536 {
191            return Err(jxl_bitstream::Error::ValidationFailed("Too many TOC entries").into());
192        }
193
194        let permutated_toc = bitstream.read_bool()?;
195        let permutation = if permutated_toc {
196            let mut decoder = jxl_coding::Decoder::parse(bitstream, 8)?;
197            decoder.begin(bitstream)?;
198            let permutation =
199                jxl_coding::read_permutation(bitstream, &mut decoder, entry_count, 0)?;
200            decoder.finalize()?;
201            permutation
202        } else {
203            Vec::new()
204        };
205
206        bitstream.zero_pad_to_byte()?;
207        let sizes = (0..entry_count)
208            .map(|_| bitstream.read_u32(U(10), 1024 + U(14), 17408 + U(22), 4211712 + U(30)))
209            .collect::<std::result::Result<Vec<_>, _>>()?;
210        bitstream.zero_pad_to_byte()?;
211
212        let mut offsets = Vec::with_capacity(sizes.len());
213        let mut acc = bitstream.num_read_bits() / 8;
214        let mut total_size = 0usize;
215        for &size in &sizes {
216            offsets.push(acc);
217            acc += size as usize;
218            total_size += size as usize;
219        }
220
221        let section_kinds = if entry_count == 1 {
222            vec![TocGroupKind::All]
223        } else {
224            let mut out = Vec::with_capacity(entry_count as usize);
225            out.push(TocGroupKind::LfGlobal);
226            for idx in 0..ctx.num_lf_groups() {
227                out.push(TocGroupKind::LfGroup(idx));
228            }
229            out.push(TocGroupKind::HfGlobal);
230            for pass_idx in 0..num_passes {
231                for group_idx in 0..num_groups {
232                    out.push(TocGroupKind::GroupPass {
233                        pass_idx,
234                        group_idx,
235                    });
236                }
237            }
238            out
239        };
240
241        let (offsets, sizes, bitstream_to_original, original_to_bitstream) = if permutated_toc {
242            let mut bitstream_to_original = vec![0usize; permutation.len()];
243            let mut offsets_out = Vec::with_capacity(permutation.len());
244            let mut sizes_out = Vec::with_capacity(permutation.len());
245            for (idx, &perm) in permutation.iter().enumerate() {
246                offsets_out.push(offsets[perm]);
247                sizes_out.push(sizes[perm]);
248                bitstream_to_original[perm] = idx;
249            }
250            (offsets_out, sizes_out, bitstream_to_original, permutation)
251        } else {
252            (offsets, sizes, Vec::new(), Vec::new())
253        };
254
255        let groups = sizes
256            .into_iter()
257            .zip(offsets)
258            .zip(section_kinds)
259            .map(|((size, offset), kind)| TocGroup { kind, offset, size })
260            .collect::<Vec<_>>();
261
262        Ok(Self {
263            num_lf_groups: ctx.num_lf_groups() as usize,
264            num_groups: num_groups as usize,
265            groups,
266            bitstream_to_original,
267            original_to_bitstream,
268            total_size,
269        })
270    }
271}
jxl_frame/data/toc.rs

jxl_frame/data/
toc.rs
