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() => panic!(
120                "Cannot request group type of {:?} for single-group frame",
121                kind
122            ),
123            TocGroupKind::All => panic!("Cannot request group type of All for multi-group frame"),
124            TocGroupKind::LfGlobal => 0,
125            TocGroupKind::LfGroup(lf_group_idx) => 1 + lf_group_idx as usize,
126            TocGroupKind::HfGlobal => 1 + self.num_lf_groups,
127            TocGroupKind::GroupPass {
128                pass_idx,
129                group_idx,
130            } => {
131                1 + self.num_lf_groups
132                    + 1
133                    + pass_idx as usize * self.num_groups
134                    + group_idx as usize
135            }
136        };
137
138        if self.original_to_bitstream.is_empty() {
139            original_order
140        } else {
141            self.original_to_bitstream[original_order]
142        }
143    }
144
145    /// Returns the total size of the frame data in bytes.
146    pub fn total_byte_size(&self) -> usize {
147        self.total_size
148    }
149
150    pub fn iter_bitstream_order(&self) -> impl Iterator<Item = TocGroup> + Send {
151        let groups = if self.bitstream_to_original.is_empty() {
152            self.groups.clone()
153        } else {
154            self.bitstream_to_original
155                .iter()
156                .map(|&idx| self.groups[idx])
157                .collect()
158        };
159        groups.into_iter()
160    }
161}
162
163impl Toc {
164    pub(crate) fn adjust_offsets(&mut self, global_frame_offset: usize) {
165        if global_frame_offset == 0 {
166            return;
167        }
168
169        for group in &mut self.groups {
170            group.offset = group
171                .offset
172                .checked_sub(global_frame_offset)
173                .expect("group offset is smaller than global frame offset");
174        }
175    }
176}
177
178impl Bundle<&crate::FrameHeader> for Toc {
179    type Error = crate::Error;
180
181    fn parse(bitstream: &mut Bitstream, ctx: &crate::FrameHeader) -> Result<Self> {
182        let num_groups = ctx.num_groups();
183        let num_passes = ctx.passes.num_passes;
184
185        let entry_count = if num_groups == 1 && num_passes == 1 {
186            1
187        } else {
188            1 + ctx.num_lf_groups() + 1 + num_groups * num_passes
189        };
190
191        if entry_count > 65536 {
192            return Err(jxl_bitstream::Error::ValidationFailed("Too many TOC entries").into());
193        }
194
195        let permutated_toc = bitstream.read_bool()?;
196        let permutation = if permutated_toc {
197            let mut decoder = jxl_coding::Decoder::parse(bitstream, 8)?;
198            decoder.begin(bitstream)?;
199            let permutation =
200                jxl_coding::read_permutation(bitstream, &mut decoder, entry_count, 0)?;
201            decoder.finalize()?;
202            permutation
203        } else {
204            Vec::new()
205        };
206
207        bitstream.zero_pad_to_byte()?;
208        let sizes = (0..entry_count)
209            .map(|_| bitstream.read_u32(U(10), 1024 + U(14), 17408 + U(22), 4211712 + U(30)))
210            .collect::<std::result::Result<Vec<_>, _>>()?;
211        bitstream.zero_pad_to_byte()?;
212
213        let mut offsets = Vec::with_capacity(sizes.len());
214        let mut acc = bitstream.num_read_bits() / 8;
215        let mut total_size = 0usize;
216        for &size in &sizes {
217            offsets.push(acc);
218            acc += size as usize;
219            total_size += size as usize;
220        }
221
222        let section_kinds = if entry_count == 1 {
223            vec![TocGroupKind::All]
224        } else {
225            let mut out = Vec::with_capacity(entry_count as usize);
226            out.push(TocGroupKind::LfGlobal);
227            for idx in 0..ctx.num_lf_groups() {
228                out.push(TocGroupKind::LfGroup(idx));
229            }
230            out.push(TocGroupKind::HfGlobal);
231            for pass_idx in 0..num_passes {
232                for group_idx in 0..num_groups {
233                    out.push(TocGroupKind::GroupPass {
234                        pass_idx,
235                        group_idx,
236                    });
237                }
238            }
239            out
240        };
241
242        let (offsets, sizes, bitstream_to_original, original_to_bitstream) = if permutated_toc {
243            let mut bitstream_to_original = vec![0usize; permutation.len()];
244            let mut offsets_out = Vec::with_capacity(permutation.len());
245            let mut sizes_out = Vec::with_capacity(permutation.len());
246            for (idx, &perm) in permutation.iter().enumerate() {
247                offsets_out.push(offsets[perm]);
248                sizes_out.push(sizes[perm]);
249                bitstream_to_original[perm] = idx;
250            }
251            (offsets_out, sizes_out, bitstream_to_original, permutation)
252        } else {
253            (offsets, sizes, Vec::new(), Vec::new())
254        };
255
256        let groups = sizes
257            .into_iter()
258            .zip(offsets)
259            .zip(section_kinds)
260            .map(|((size, offset), kind)| TocGroup { kind, offset, size })
261            .collect::<Vec<_>>();
262
263        Ok(Self {
264            num_lf_groups: ctx.num_lf_groups() as usize,
265            num_groups: num_groups as usize,
266            groups,
267            bitstream_to_original,
268            original_to_bitstream,
269            total_size,
270        })
271    }
272}
jxl_frame/data/toc.rs

jxl_frame/data/
toc.rs
