j2k_jpeg/adapter/
device_plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use crate::allocation::{
4    checked_add_allocation_bytes, checked_allocation_bytes, ensure_allocation_bytes,
5    try_reserve_for_len_with_live_budget,
6};
7use crate::decoder::Decoder;
8use crate::error::JpegError;
9use crate::info::{ColorSpace, SofKind};
10use crate::internal::checkpoint::{
11    build_checkpoint_plan_from_validated_with_live_budget, checkpoint_count_summary, total_mcus,
12    validate_scan_bytes, DeviceCheckpoint,
13};
14use crate::Warning;
15use alloc::borrow::Cow;
16use alloc::vec::Vec;
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19#[doc(hidden)]
20/// Component metadata needed by device-side JPEG decoders.
21pub struct DeviceComponentPlan {
22    /// Horizontal sampling factor.
23    pub h: u8,
24    /// Vertical sampling factor.
25    pub v: u8,
26    /// Output component index in decoded pixel/component order.
27    pub output_index: usize,
28}
29
30#[derive(Debug, PartialEq, Eq)]
31#[doc(hidden)]
32/// Validated decode plan and entropy payload for device-side decoders.
33///
34/// The plan is move-only because an owned scan payload, warning list, and
35/// checkpoint graph can together approach the shared host-allocation cap.
36/// Borrow it or use `Arc` when explicit shared ownership is required.
37pub struct DeviceDecodePlan<'a> {
38    /// Image dimensions as `(width, height)` in pixels.
39    pub dimensions: (u32, u32),
40    /// Header-derived JPEG color space.
41    pub color_space: ColorSpace,
42    /// Restart interval in MCUs, if restart markers are present.
43    pub restart_interval: Option<u16>,
44    /// Non-fatal warnings collected while building the plan.
45    pub warnings: Vec<Warning>,
46    /// Entropy scan bytes with device-compatible marker handling.
47    pub scan_bytes: Cow<'a, [u8]>,
48    /// Component sampling and output placement.
49    pub components: Vec<DeviceComponentPlan>,
50    /// Entropy checkpoints for restart or synthetic resume points.
51    pub checkpoints: Vec<DeviceCheckpoint>,
52    /// True when the tile matches the Metal/CUDA fast 4:2:0 path.
53    pub matches_fast_420: bool,
54    /// True when the tile matches the Metal/CUDA fast 4:2:2 path.
55    pub matches_fast_422: bool,
56    /// True when the tile matches the Metal/CUDA fast 4:4:4 path.
57    pub matches_fast_444: bool,
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61/// Lightweight summary used for batching without materializing scan bytes.
62pub struct DeviceBatchSummary {
63    /// Restart interval in MCUs, if restart markers are present.
64    pub restart_interval: Option<u16>,
65    /// Number of checkpoints the full device plan would contain.
66    pub checkpoint_count: usize,
67    /// True when the tile matches the fast 4:2:0 device path.
68    pub matches_fast_420: bool,
69    /// True when the tile matches the fast 4:2:2 device path.
70    pub matches_fast_422: bool,
71    /// True when the tile matches the fast 4:4:4 device path.
72    pub matches_fast_444: bool,
73}
74
75/// Build a full device decode plan for an inspected decoder.
76#[doc(hidden)]
77pub fn build_device_plan<'a>(
78    decoder: &'a Decoder<'a>,
79    cadence_mcus: u32,
80) -> Result<DeviceDecodePlan<'a>, JpegError> {
81    if !matches!(
82        decoder.info().sof_kind,
83        SofKind::Baseline8 | SofKind::Extended8
84    ) {
85        return Err(JpegError::NotImplemented {
86            sof: decoder.info().sof_kind,
87        });
88    }
89    let plan = &decoder.plan;
90    let restart_interval = plan.restart_interval.filter(|&interval| interval > 0);
91    let validated_scan = validate_scan_bytes(
92        &decoder.bytes[plan.scan_offset..],
93        restart_interval.is_some(),
94        plan.scan_offset,
95    )?;
96    let scan_bytes = Cow::Borrowed(validated_scan.payload());
97    let missing_eoi = validated_scan.is_missing_eoi();
98    let warning_count = decoder
99        .warnings
100        .len()
101        .checked_add(usize::from(missing_eoi))
102        .ok_or(JpegError::MemoryCapExceeded {
103            requested: usize::MAX,
104            cap: j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES,
105        })?;
106    let expected_checkpoint_count = checkpoint_count_summary(
107        total_mcus(plan),
108        cadence_mcus.max(1),
109        restart_interval.map(u32::from),
110    );
111    let retained_decoder_bytes = retained_decoder_allocation_bytes(decoder)?;
112    let output_bytes = device_plan_output_allocation_bytes(
113        expected_checkpoint_count,
114        warning_count,
115        plan.components.len(),
116    )?;
117    let terminated_copy_bytes =
118        validated_scan.terminated_copy_len(j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES)?;
119    ensure_allocation_bytes(checked_add_allocation_bytes(
120        checked_add_allocation_bytes(retained_decoder_bytes, output_bytes)?,
121        terminated_copy_bytes,
122    )?)?;
123
124    let mut live_bytes = retained_decoder_bytes;
125    let mut warnings = Vec::new();
126    try_reserve_for_len_with_live_budget(
127        &mut warnings,
128        warning_count,
129        &mut live_bytes,
130        j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES,
131    )?;
132    warnings.extend_from_slice(&decoder.warnings);
133    if missing_eoi {
134        warnings.push(Warning::MissingEoi);
135    }
136
137    let mut components = Vec::new();
138    try_reserve_for_len_with_live_budget(
139        &mut components,
140        plan.components.len(),
141        &mut live_bytes,
142        j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES,
143    )?;
144    components.extend(plan.components.iter().map(|component| DeviceComponentPlan {
145        h: component.h,
146        v: component.v,
147        output_index: component.output_index,
148    }));
149    let checkpoints = build_checkpoint_plan_from_validated_with_live_budget(
150        plan,
151        validated_scan,
152        cadence_mcus,
153        &mut live_bytes,
154        j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES,
155    )?;
156    if checkpoints.len() != expected_checkpoint_count {
157        return Err(JpegError::InternalInvariant {
158            reason: "device checkpoint count disagrees with its allocation plan",
159        });
160    }
161
162    Ok(DeviceDecodePlan {
163        dimensions: plan.dimensions,
164        color_space: plan.color_space,
165        restart_interval,
166        warnings,
167        scan_bytes,
168        components,
169        checkpoints,
170        matches_fast_420: plan.matches_fast_tile_shape(),
171        matches_fast_422: plan.matches_fast_rgb422_shape(),
172        matches_fast_444: plan.matches_fast_rgb444_shape(),
173    })
174}
175
176pub(super) fn retained_decoder_allocation_bytes(decoder: &Decoder<'_>) -> Result<usize, JpegError> {
177    let baseline = decoder.retained_allocation_bytes_excluding_cpu_checkpoint_cache()?;
178    let checkpoint_cache_bytes = decoder
179        .cpu_entropy_checkpoints
180        .lock()
181        .map_err(|_| JpegError::InternalInvariant {
182            reason: "CPU entropy checkpoint cache mutex poisoned",
183        })?
184        .retained_allocation_bytes()?;
185    checked_add_allocation_bytes(baseline, checkpoint_cache_bytes)
186}
187
188fn device_plan_output_allocation_bytes(
189    checkpoint_count: usize,
190    warning_count: usize,
191    component_count: usize,
192) -> Result<usize, JpegError> {
193    let mut total = checked_allocation_bytes::<DeviceCheckpoint>(checkpoint_count)?;
194    total =
195        checked_add_allocation_bytes(total, checked_allocation_bytes::<Warning>(warning_count)?)?;
196    checked_add_allocation_bytes(
197        total,
198        checked_allocation_bytes::<DeviceComponentPlan>(component_count)?,
199    )
200}
201
202/// Summarize device planning properties without copying entropy data.
203#[doc(hidden)]
204pub fn summarize_device_batch(decoder: &Decoder<'_>, cadence_mcus: u32) -> DeviceBatchSummary {
205    if !matches!(
206        decoder.info().sof_kind,
207        SofKind::Baseline8 | SofKind::Extended8
208    ) {
209        return DeviceBatchSummary {
210            restart_interval: None,
211            checkpoint_count: 0,
212            matches_fast_420: false,
213            matches_fast_422: false,
214            matches_fast_444: false,
215        };
216    }
217    let plan = &decoder.plan;
218    let restart_interval = plan.restart_interval.filter(|&interval| interval > 0);
219    let checkpoint_count = checkpoint_count_summary(
220        total_mcus(plan),
221        cadence_mcus,
222        restart_interval.map(u32::from),
223    );
224
225    DeviceBatchSummary {
226        restart_interval,
227        checkpoint_count,
228        matches_fast_420: plan.matches_fast_tile_shape(),
229        matches_fast_422: plan.matches_fast_rgb422_shape(),
230        matches_fast_444: plan.matches_fast_rgb444_shape(),
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237
238    #[test]
239    fn device_plan_output_rejects_aggregate_retained_vectors() {
240        let cap = j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES;
241        let checkpoint_count = (cap * 3 / 5) / core::mem::size_of::<DeviceCheckpoint>();
242        let warning_count = (cap * 3 / 5) / core::mem::size_of::<Warning>();
243        assert!(checked_allocation_bytes::<DeviceCheckpoint>(checkpoint_count).is_ok());
244        assert!(checked_allocation_bytes::<Warning>(warning_count).is_ok());
245        assert!(matches!(
246            device_plan_output_allocation_bytes(checkpoint_count, warning_count, 0),
247            Err(JpegError::MemoryCapExceeded { requested, cap: limit })
248                if requested > limit && limit == cap
249        ));
250    }
251
252    #[test]
253    fn device_plan_output_boundary_counts_all_three_payloads() {
254        let checkpoints = 3usize;
255        let warnings = 2usize;
256        let components = 4usize;
257        let expected = checkpoints * core::mem::size_of::<DeviceCheckpoint>()
258            + warnings * core::mem::size_of::<Warning>()
259            + components * core::mem::size_of::<DeviceComponentPlan>();
260        assert_eq!(
261            device_plan_output_allocation_bytes(checkpoints, warnings, components).unwrap(),
262            expected
263        );
264    }
265
266    #[test]
267    fn retained_decoder_bytes_include_populated_cpu_checkpoint_cache() {
268        let bytes = j2k_test_support::minimal_baseline_jpeg();
269        let decoder = Decoder::new(&bytes).unwrap();
270        let before = retained_decoder_allocation_bytes(&decoder).unwrap();
271        let retained_checkpoint_bytes = {
272            let mut cache = decoder.cpu_entropy_checkpoints.lock().unwrap();
273            cache.checkpoints.try_reserve_exact(4).unwrap();
274            cache.retained_allocation_bytes().unwrap()
275        };
276        let after = retained_decoder_allocation_bytes(&decoder).unwrap();
277
278        assert!(retained_checkpoint_bytes >= 4 * core::mem::size_of::<DeviceCheckpoint>());
279        assert_eq!(after, before + retained_checkpoint_bytes);
280    }
281}
j2k_jpeg/adapter/device_plan.rs

j2k_jpeg/adapter/
device_plan.rs
