Skip to main content

jxr_core/
plan.rs

1//! Parsed reconstruction plans and macroblock-major coefficient storage.
2
3use 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/// Quantization steps selected for one macroblock.
19#[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/// Tile and image boundary state consumed by overlap kernels.
27#[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/// Structure-of-arrays metadata for GPU and CPU reconstruction.
63#[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    /// HP prediction direction selected from reconstructed DC/LP coefficients.
70    pub hp_predictions: Vec<PredictionMode>,
71    pub tile_edges: Vec<TileEdgeFlags>,
72    /// Macroblock column in the extended coded plane.
73    pub coded_x: Vec<u32>,
74    /// Macroblock row in the extended coded plane.
75    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/// Location of a plane within a coefficient arena.
123#[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    /// Number of 4-sample blocks across one component macroblock.
130    pub block_columns: u8,
131    /// Number of 4-sample blocks down one component macroblock.
132    pub block_rows: u8,
133}
134
135/// Macroblock-major signed coefficients plus their reconstruction metadata.
136#[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/// Coefficient geometry and metadata for externally owned coefficient storage.
144#[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/// Per-plane reconstruction geometry and boundary policy.
209#[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/// Independently located compressed tile and its affected output region.
220#[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    /// Whether this tile intersects the requested reconstruction window.
228    pub required_for_reconstruction: bool,
229}
230
231/// Device-neutral result of parse and resource planning.
232#[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    /// Requested region in full-resolution source coordinates.
241    pub output_region: Rect,
242    /// Actual output region in the selected native-resolution coordinate grid.
243    pub decoded_region: Rect,
244    pub scale: DecodeScale,
245    pub coefficient_bytes: usize,
246}
247
248impl PreparedPlan {
249    /// Algorithmic reconstruction work after sparse entropy coefficients have
250    /// been expanded into component sample blocks.
251    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}