j2k_types/decode_plan/
allocation.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use core::mem::size_of;
4
5use super::{
6    HtOwnedCodeBlockBatchJob, J2kDirectColorPlan, J2kDirectGrayscalePlan, J2kDirectGrayscaleStep,
7    J2kDirectRgbaPlan, J2kOwnedCodeBlockBatchJob, J2kReferencedClassicPlan, J2kReferencedHtj2kPlan,
8    J2kReferencedTilePlan, DEFAULT_MAX_DECODE_BYTES,
9};
10use crate::{
11    HtCodeBlockPayloadRanges, J2kClassicCodeBlockPayload, J2kCodeBlockSegment, J2kCodestreamRange,
12};
13
14/// Retained decode-plan capacity overflow or limit violation.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct DecodePlanAllocationError;
17
18#[derive(Default)]
19struct Budget {
20    bytes: usize,
21}
22
23impl Budget {
24    fn include_capacity<T>(&mut self, capacity: usize) -> Result<(), DecodePlanAllocationError> {
25        let bytes = capacity
26            .checked_mul(size_of::<T>())
27            .ok_or(DecodePlanAllocationError)?;
28        self.bytes = self
29            .bytes
30            .checked_add(bytes)
31            .ok_or(DecodePlanAllocationError)?;
32        (self.bytes <= DEFAULT_MAX_DECODE_BYTES)
33            .then_some(())
34            .ok_or(DecodePlanAllocationError)
35    }
36}
37
38impl J2kDirectGrayscalePlan {
39    /// Return allocator capacities retained by this plan's owner graph.
40    pub fn retained_allocation_bytes(&self) -> Result<usize, DecodePlanAllocationError> {
41        let mut budget = Budget::default();
42        include_grayscale(&mut budget, self)?;
43        Ok(budget.bytes)
44    }
45}
46
47impl J2kDirectColorPlan {
48    /// Return allocator capacities retained by this plan's owner graph.
49    pub fn retained_allocation_bytes(&self) -> Result<usize, DecodePlanAllocationError> {
50        retained_components(&self.component_plans, self.component_plans.capacity())
51    }
52}
53
54impl J2kDirectRgbaPlan {
55    /// Return allocator capacities retained by this plan's owner graph.
56    pub fn retained_allocation_bytes(&self) -> Result<usize, DecodePlanAllocationError> {
57        retained_components(&self.component_plans, self.component_plans.capacity())
58    }
59}
60
61impl J2kReferencedHtj2kPlan {
62    /// Return allocator capacities retained by geometry and payload ranges.
63    pub fn retained_allocation_bytes(&self) -> Result<usize, DecodePlanAllocationError> {
64        let mut budget = Budget::default();
65        match self {
66            Self::Grayscale {
67                tiles, payloads, ..
68            }
69            | Self::Color {
70                tiles, payloads, ..
71            }
72            | Self::Rgba {
73                tiles, payloads, ..
74            } => {
75                include_tiles(&mut budget, tiles, tiles.capacity())?;
76                budget.include_capacity::<HtCodeBlockPayloadRanges>(payloads.capacity())?;
77            }
78        }
79        Ok(budget.bytes)
80    }
81}
82
83impl J2kReferencedClassicPlan {
84    /// Return allocator capacities retained by geometry and payload ranges.
85    pub fn retained_allocation_bytes(&self) -> Result<usize, DecodePlanAllocationError> {
86        let mut budget = Budget::default();
87        match self {
88            Self::Grayscale {
89                tiles,
90                payloads,
91                ranges,
92                ..
93            }
94            | Self::Color {
95                tiles,
96                payloads,
97                ranges,
98                ..
99            }
100            | Self::Rgba {
101                tiles,
102                payloads,
103                ranges,
104                ..
105            } => {
106                include_tiles(&mut budget, tiles, tiles.capacity())?;
107                include_classic(&mut budget, payloads.capacity(), ranges.capacity())?;
108            }
109        }
110        Ok(budget.bytes)
111    }
112}
113
114fn include_tiles(
115    budget: &mut Budget,
116    tiles: &[J2kReferencedTilePlan],
117    capacity: usize,
118) -> Result<(), DecodePlanAllocationError> {
119    budget.include_capacity::<J2kReferencedTilePlan>(capacity)?;
120    for tile in tiles {
121        include_classic(
122            budget,
123            tile.classic_payloads.capacity(),
124            tile.classic_ranges.capacity(),
125        )?;
126        if let Some(plan) = tile.grayscale_geometry() {
127            include_grayscale(budget, plan)?;
128        } else if let Some(plan) = tile.color_geometry() {
129            include_components(
130                budget,
131                &plan.component_plans,
132                plan.component_plans.capacity(),
133            )?;
134        } else if let Some(plan) = tile.rgba_geometry() {
135            include_components(
136                budget,
137                &plan.component_plans,
138                plan.component_plans.capacity(),
139            )?;
140        } else {
141            return Err(DecodePlanAllocationError);
142        }
143    }
144    Ok(())
145}
146
147fn include_classic(
148    budget: &mut Budget,
149    payload_capacity: usize,
150    range_capacity: usize,
151) -> Result<(), DecodePlanAllocationError> {
152    budget.include_capacity::<J2kClassicCodeBlockPayload>(payload_capacity)?;
153    budget.include_capacity::<J2kCodestreamRange>(range_capacity)
154}
155
156fn retained_components(
157    components: &[J2kDirectGrayscalePlan],
158    capacity: usize,
159) -> Result<usize, DecodePlanAllocationError> {
160    let mut budget = Budget::default();
161    include_components(&mut budget, components, capacity)?;
162    Ok(budget.bytes)
163}
164
165fn include_components(
166    budget: &mut Budget,
167    components: &[J2kDirectGrayscalePlan],
168    capacity: usize,
169) -> Result<(), DecodePlanAllocationError> {
170    budget.include_capacity::<J2kDirectGrayscalePlan>(capacity)?;
171    for component in components {
172        include_grayscale(budget, component)?;
173    }
174    Ok(())
175}
176
177fn include_grayscale(
178    budget: &mut Budget,
179    plan: &J2kDirectGrayscalePlan,
180) -> Result<(), DecodePlanAllocationError> {
181    budget.include_capacity::<J2kDirectGrayscaleStep>(plan.steps.capacity())?;
182    for step in &plan.steps {
183        match step {
184            J2kDirectGrayscaleStep::ClassicSubBand(subband) => {
185                budget.include_capacity::<J2kOwnedCodeBlockBatchJob>(subband.jobs.capacity())?;
186                for job in &subband.jobs {
187                    budget.include_capacity::<u8>(job.data.capacity())?;
188                    budget.include_capacity::<J2kCodeBlockSegment>(job.segments.capacity())?;
189                }
190            }
191            J2kDirectGrayscaleStep::HtSubBand(subband) => {
192                budget.include_capacity::<HtOwnedCodeBlockBatchJob>(subband.jobs.capacity())?;
193                for job in &subband.jobs {
194                    budget.include_capacity::<u8>(job.data.capacity())?;
195                }
196            }
197            J2kDirectGrayscaleStep::Idwt(_) | J2kDirectGrayscaleStep::Store(_) => {}
198        }
199    }
200    Ok(())
201}
202
203#[cfg(test)]
204mod tests {
205    use alloc::vec::Vec;
206
207    use super::*;
208    use crate::{J2kCodeBlockStyle, J2kSubBandType};
209
210    #[test]
211    fn retained_bytes_use_nested_vector_capacities() {
212        let mut jobs = Vec::new();
213        jobs.try_reserve_exact(3).expect("job capacity");
214        jobs.push(J2kOwnedCodeBlockBatchJob {
215            output_x: 0,
216            output_y: 0,
217            data: {
218                let mut data = Vec::new();
219                data.try_reserve_exact(11).expect("data capacity");
220                data
221            },
222            segments: {
223                let mut segments = Vec::new();
224                segments.try_reserve_exact(5).expect("segment capacity");
225                segments
226            },
227            width: 1,
228            height: 1,
229            output_stride: 1,
230            missing_bit_planes: 0,
231            number_of_coding_passes: 0,
232            total_bitplanes: 0,
233            roi_shift: 0,
234            sub_band_type: J2kSubBandType::LowLow,
235            style: J2kCodeBlockStyle {
236                selective_arithmetic_coding_bypass: false,
237                reset_context_probabilities: false,
238                termination_on_each_pass: false,
239                vertically_causal_context: false,
240                segmentation_symbols: false,
241            },
242            strict: false,
243            dequantization_step: 1.0,
244        });
245        let mut steps = Vec::new();
246        steps.try_reserve_exact(4).expect("step capacity");
247        steps.push(J2kDirectGrayscaleStep::ClassicSubBand(
248            super::super::J2kOwnedSubBandPlan {
249                band_id: 0,
250                rect: super::super::J2kRect {
251                    x0: 0,
252                    y0: 0,
253                    x1: 1,
254                    y1: 1,
255                },
256                width: 1,
257                height: 1,
258                irreversible_midpoint: false,
259                jobs,
260            },
261        ));
262        let plan = J2kDirectGrayscalePlan {
263            dimensions: (1, 1),
264            bit_depth: 8,
265            steps,
266        };
267
268        assert_eq!(
269            plan.retained_allocation_bytes().expect("retained bytes"),
270            4 * size_of::<J2kDirectGrayscaleStep>()
271                + 3 * size_of::<J2kOwnedCodeBlockBatchJob>()
272                + 11
273                + 5 * size_of::<J2kCodeBlockSegment>()
274        );
275    }
276
277    #[test]
278    fn direct_plan_owner_types_remain_move_only_values() {
279        fn assert_debug<T: core::fmt::Debug>() {}
280
281        assert_debug::<J2kDirectColorPlan>();
282        assert_debug::<J2kDirectGrayscalePlan>();
283        assert_debug::<super::super::J2kOwnedSubBandPlan>();
284        assert_debug::<super::super::HtOwnedSubBandPlan>();
285        assert_debug::<J2kOwnedCodeBlockBatchJob>();
286        assert_debug::<HtOwnedCodeBlockBatchJob>();
287    }
288}
j2k_types/decode_plan/allocation.rs

j2k_types/decode_plan/
allocation.rs
