j2k-native 0.7.0

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Retained-allocation accounting for native direct-plan owner graphs.

use core::mem::size_of;

use crate::{
    HtOwnedCodeBlockBatchJob, J2kCodeBlockSegment, J2kDirectColorPlan, J2kDirectGrayscalePlan,
    J2kDirectGrayscaleStep, J2kOwnedCodeBlockBatchJob, Result, ValidationError,
    DEFAULT_MAX_DECODE_BYTES,
};

#[derive(Default)]
struct RetainedPlanBudget {
    bytes: usize,
}

impl RetainedPlanBudget {
    fn include_capacity<T>(&mut self, capacity: usize) -> Result<()> {
        let bytes = capacity
            .checked_mul(size_of::<T>())
            .ok_or(ValidationError::ImageTooLarge)?;
        self.include_bytes(bytes)
    }

    fn include_bytes(&mut self, bytes: usize) -> Result<()> {
        self.bytes = self
            .bytes
            .checked_add(bytes)
            .ok_or(ValidationError::ImageTooLarge)?;
        if self.bytes > DEFAULT_MAX_DECODE_BYTES {
            return Err(ValidationError::ImageTooLarge.into());
        }
        Ok(())
    }
}

impl J2kDirectGrayscalePlan {
    /// Return the allocator capacities retained by this direct-plan owner graph.
    ///
    /// The root plan value itself is not included. Callers that place the root
    /// in a separate heap allocation must account that allocation separately.
    ///
    /// # Errors
    ///
    /// Returns an error if nested capacity arithmetic overflows or exceeds the
    /// native decode allocation ceiling.
    #[doc(hidden)]
    pub fn retained_allocation_bytes(&self) -> Result<usize> {
        let mut budget = RetainedPlanBudget::default();
        include_grayscale_plan(&mut budget, self)?;
        Ok(budget.bytes)
    }
}

impl J2kDirectColorPlan {
    /// Return the allocator capacities retained by this direct-plan owner graph.
    ///
    /// The root plan value itself is not included. Callers that place the root
    /// in a separate heap allocation must account that allocation separately.
    ///
    /// # Errors
    ///
    /// Returns an error if nested capacity arithmetic overflows or exceeds the
    /// native decode allocation ceiling.
    #[doc(hidden)]
    pub fn retained_allocation_bytes(&self) -> Result<usize> {
        let mut budget = RetainedPlanBudget::default();
        budget.include_capacity::<J2kDirectGrayscalePlan>(self.component_plans.capacity())?;
        for component in &self.component_plans {
            include_grayscale_plan(&mut budget, component)?;
        }
        Ok(budget.bytes)
    }
}

fn include_grayscale_plan(
    budget: &mut RetainedPlanBudget,
    plan: &J2kDirectGrayscalePlan,
) -> Result<()> {
    budget.include_capacity::<J2kDirectGrayscaleStep>(plan.steps.capacity())?;
    for step in &plan.steps {
        match step {
            J2kDirectGrayscaleStep::ClassicSubBand(sub_band) => {
                budget.include_capacity::<J2kOwnedCodeBlockBatchJob>(sub_band.jobs.capacity())?;
                for job in &sub_band.jobs {
                    budget.include_capacity::<u8>(job.data.capacity())?;
                    budget.include_capacity::<J2kCodeBlockSegment>(job.segments.capacity())?;
                }
            }
            J2kDirectGrayscaleStep::HtSubBand(sub_band) => {
                budget.include_capacity::<HtOwnedCodeBlockBatchJob>(sub_band.jobs.capacity())?;
                for job in &sub_band.jobs {
                    budget.include_capacity::<u8>(job.data.capacity())?;
                }
            }
            J2kDirectGrayscaleStep::Idwt(_) | J2kDirectGrayscaleStep::Store(_) => {}
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use alloc::vec::Vec;

    use super::*;
    use crate::{HtOwnedSubBandPlan, J2kOwnedSubBandPlan, J2kRect};

    #[test]
    fn retained_bytes_use_nested_vector_capacities() {
        let mut jobs = Vec::new();
        jobs.try_reserve_exact(3).expect("job test capacity");
        jobs.push(J2kOwnedCodeBlockBatchJob {
            output_x: 0,
            output_y: 0,
            data: {
                let mut data = Vec::new();
                data.try_reserve_exact(11).expect("data test capacity");
                data
            },
            segments: {
                let mut segments = Vec::new();
                segments
                    .try_reserve_exact(5)
                    .expect("segment test capacity");
                segments
            },
            width: 1,
            height: 1,
            output_stride: 1,
            missing_bit_planes: 0,
            number_of_coding_passes: 0,
            total_bitplanes: 0,
            roi_shift: 0,
            sub_band_type: crate::J2kSubBandType::LowLow,
            style: crate::J2kCodeBlockStyle {
                selective_arithmetic_coding_bypass: false,
                reset_context_probabilities: false,
                termination_on_each_pass: false,
                vertically_causal_context: false,
                segmentation_symbols: false,
            },
            strict: false,
            dequantization_step: 1.0,
        });
        let mut steps = Vec::new();
        steps.try_reserve_exact(4).expect("step test capacity");
        steps.push(J2kDirectGrayscaleStep::ClassicSubBand(
            J2kOwnedSubBandPlan {
                band_id: 0,
                rect: J2kRect {
                    x0: 0,
                    y0: 0,
                    x1: 1,
                    y1: 1,
                },
                width: 1,
                height: 1,
                jobs,
            },
        ));
        let plan = J2kDirectGrayscalePlan {
            dimensions: (1, 1),
            bit_depth: 8,
            steps,
        };

        assert_eq!(
            plan.retained_allocation_bytes().expect("retained bytes"),
            4 * size_of::<J2kDirectGrayscaleStep>()
                + 3 * size_of::<J2kOwnedCodeBlockBatchJob>()
                + 11
                + 5 * size_of::<J2kCodeBlockSegment>()
        );
    }

    #[test]
    fn all_direct_plan_owner_types_are_move_only() {
        fn assert_debug<T: core::fmt::Debug>() {}

        assert_debug::<J2kDirectColorPlan>();
        assert_debug::<J2kDirectGrayscalePlan>();
        assert_debug::<J2kOwnedSubBandPlan>();
        assert_debug::<HtOwnedSubBandPlan>();
        assert_debug::<J2kOwnedCodeBlockBatchJob>();
        assert_debug::<HtOwnedCodeBlockBatchJob>();
    }
}