1use alloc::vec::Vec;
4
5use crate::{
6 AlphaMode, BandPresence, ByteRange, ChromaSampling, ColorFormat, DecodeScale, ImageInfo,
7 JxrError, JxrErrorKind, OverlapMode, Rect,
8};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub enum PredictionMode {
12 None,
13 FromLeft,
14 FromTop,
15 FromTopLeft,
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub struct QuantizerSet {
21 pub dc: u32,
22 pub low_pass: u32,
23 pub high_pass: u32,
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
28pub struct TileEdgeFlags(u8);
29
30impl TileEdgeFlags {
31 pub const LEFT: Self = Self(1 << 0);
32 pub const TOP: Self = Self(1 << 1);
33 pub const RIGHT: Self = Self(1 << 2);
34 pub const BOTTOM: Self = Self(1 << 3);
35 pub const HARD_TILE: Self = Self(1 << 4);
36
37 #[must_use]
38 pub const fn from_bits(bits: u8) -> Option<Self> {
39 if bits & !0x1f == 0 {
40 Some(Self(bits))
41 } else {
42 None
43 }
44 }
45
46 #[must_use]
47 pub const fn bits(self) -> u8 {
48 self.0
49 }
50
51 #[must_use]
52 pub const fn contains(self, flag: Self) -> bool {
53 self.0 & flag.0 == flag.0
54 }
55
56 #[must_use]
57 pub const fn union(self, flag: Self) -> Self {
58 Self(self.0 | flag.0)
59 }
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, Default)]
64pub struct MacroblockMetadata {
65 pub coefficient_offsets: Vec<u32>,
66 pub quantizers: Vec<QuantizerSet>,
67 pub bands: Vec<BandPresence>,
68 pub predictions: Vec<PredictionMode>,
69 pub hp_predictions: Vec<PredictionMode>,
71 pub tile_edges: Vec<TileEdgeFlags>,
72 pub coded_x: Vec<u32>,
74 pub coded_y: Vec<u32>,
76 pub output_x: Vec<u32>,
77 pub output_y: Vec<u32>,
78}
79
80impl MacroblockMetadata {
81 #[must_use]
82 pub fn len(&self) -> usize {
83 self.coefficient_offsets.len()
84 }
85
86 #[must_use]
87 pub fn is_empty(&self) -> bool {
88 self.coefficient_offsets.is_empty()
89 }
90
91 pub fn validate(&self, coefficient_count: usize) -> Result<(), JxrError> {
92 let expected = self.len();
93 let lengths = [
94 self.quantizers.len(),
95 self.bands.len(),
96 self.predictions.len(),
97 self.hp_predictions.len(),
98 self.tile_edges.len(),
99 self.coded_x.len(),
100 self.coded_y.len(),
101 self.output_x.len(),
102 self.output_y.len(),
103 ];
104 if lengths.iter().any(|&length| length != expected) {
105 return Err(JxrError::new(
106 JxrErrorKind::InternalInvariant,
107 "macroblock metadata lengths",
108 ));
109 }
110 if self.coefficient_offsets.iter().any(|&offset| {
111 usize::try_from(offset).map_or(true, |offset| offset > coefficient_count)
112 }) {
113 return Err(JxrError::new(
114 JxrErrorKind::InvalidSyntax,
115 "coefficient offset",
116 ));
117 }
118 Ok(())
119 }
120}
121
122#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
124pub struct CoefficientPlane {
125 pub coefficient_offset: usize,
126 pub coefficient_count: usize,
127 pub macroblock_offset: usize,
128 pub macroblock_count: usize,
129 pub block_columns: u8,
131 pub block_rows: u8,
133}
134
135#[derive(Debug, Clone, PartialEq, Eq, Default)]
137pub struct CoefficientArena {
138 pub coefficients: Vec<i32>,
139 pub macroblocks: MacroblockMetadata,
140 pub planes: Vec<CoefficientPlane>,
141}
142
143#[derive(Debug, Clone, PartialEq, Eq, Default)]
145pub struct CoefficientArenaDescriptor {
146 pub coefficient_count: usize,
147 pub macroblocks: MacroblockMetadata,
148 pub planes: Vec<CoefficientPlane>,
149}
150
151impl CoefficientArenaDescriptor {
152 pub fn validate(&self) -> Result<(), JxrError> {
153 validate_coefficient_contract(self.coefficient_count, &self.macroblocks, &self.planes)
154 }
155}
156
157impl CoefficientArena {
158 pub fn validate(&self) -> Result<(), JxrError> {
159 validate_coefficient_contract(self.coefficients.len(), &self.macroblocks, &self.planes)
160 }
161}
162
163fn validate_coefficient_contract(
164 coefficient_count: usize,
165 macroblocks: &MacroblockMetadata,
166 planes: &[CoefficientPlane],
167) -> Result<(), JxrError> {
168 macroblocks.validate(coefficient_count)?;
169 for plane in planes {
170 if !(1..=4).contains(&plane.block_columns) || !(1..=4).contains(&plane.block_rows) {
171 return Err(JxrError::new(
172 JxrErrorKind::InternalInvariant,
173 "coefficient plane block geometry",
174 ));
175 }
176 check_range(
177 plane.coefficient_offset,
178 plane.coefficient_count,
179 coefficient_count,
180 "coefficient plane range",
181 )?;
182 check_range(
183 plane.macroblock_offset,
184 plane.macroblock_count,
185 macroblocks.len(),
186 "macroblock plane range",
187 )?;
188 }
189 Ok(())
190}
191
192fn check_range(
193 offset: usize,
194 length: usize,
195 available: usize,
196 operation: &'static str,
197) -> Result<(), JxrError> {
198 let end = offset
199 .checked_add(length)
200 .ok_or_else(|| JxrError::arithmetic(operation))?;
201 if end > available {
202 Err(JxrError::new(JxrErrorKind::InternalInvariant, operation))
203 } else {
204 Ok(())
205 }
206}
207
208#[derive(Debug, Clone, PartialEq, Eq)]
210pub struct PlanePlan {
211 pub width: u32,
212 pub height: u32,
213 pub macroblocks_x: u32,
214 pub macroblocks_y: u32,
215 pub overlap: OverlapMode,
216 pub coefficient_plane: usize,
217}
218
219#[derive(Debug, Clone, PartialEq, Eq)]
221pub struct TilePlan {
222 pub packet_range: ByteRange,
223 pub output_region: Rect,
224 pub macroblock_start: u32,
225 pub macroblock_count: u32,
226 pub hard_boundaries: bool,
227 pub required_for_reconstruction: bool,
229}
230
231#[derive(Debug, Clone, PartialEq, Eq)]
233pub struct PreparedPlan {
234 pub info: ImageInfo,
235 pub codestream_range: ByteRange,
236 pub primary: PlanePlan,
237 pub alpha: Option<PlanePlan>,
238 pub tiles: Vec<TilePlan>,
239 pub reconstruction_region: Rect,
240 pub output_region: Rect,
242 pub decoded_region: Rect,
244 pub scale: DecodeScale,
245 pub coefficient_bytes: usize,
246}
247
248impl PreparedPlan {
249 pub fn reconstructed_coefficients(&self) -> Result<u64, JxrError> {
252 let macroblocks = self
253 .tiles
254 .iter()
255 .filter(|tile| tile.required_for_reconstruction)
256 .try_fold(0_u64, |total, tile| {
257 total
258 .checked_add(u64::from(tile.macroblock_count))
259 .ok_or_else(|| JxrError::arithmetic("reconstruction macroblock count"))
260 })?;
261 let primary = match self.info.primary.color_format {
262 ColorFormat::Luma => 256,
263 ColorFormat::Yuv(ChromaSampling::Cs420) => 384,
264 ColorFormat::Yuv(ChromaSampling::Cs422) => 512,
265 ColorFormat::Yuv(ChromaSampling::Cs444) | ColorFormat::Rgb | ColorFormat::Rgbe => 768,
266 ColorFormat::Cmyk | ColorFormat::CmykDirect | ColorFormat::YuvK => 1_024,
267 ColorFormat::NComponent(components) => u64::from(components)
268 .checked_mul(256)
269 .ok_or_else(|| JxrError::arithmetic("N-component reconstruction work"))?,
270 };
271 let alpha = match self.info.alpha_mode {
272 AlphaMode::Integrated => 256,
273 AlphaMode::Separate if self.alpha.is_some() => 256,
274 AlphaMode::None | AlphaMode::Separate => 0,
275 };
276 macroblocks
277 .checked_mul(primary + alpha)
278 .ok_or_else(|| JxrError::arithmetic("reconstructed coefficient count"))
279 }
280}