j2k-native 0.13.1

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

//! Aggregate planning and reusable-capacity policy for decomposition storage.

use super::{CodeBlock, Decomposition, Layer, Precinct, Segment, SubBand, SCRATCH_PRECINCTS};
use crate::error::{DecodingError, Result, ValidationError};
use crate::j2c::decode::{DecompositionStorage, TileDecompositions};
use crate::j2c::rect::IntRect;
use crate::j2c::tag_tree::TagNode;
use crate::j2c::tile::Tile;
use crate::DEFAULT_MAX_DECODE_BYTES;
use core::mem::size_of;

mod plan;
mod reuse;

pub(super) use plan::{code_block_grid, tag_tree_node_count};
pub(super) use reuse::push_preallocated;

/// How sub-band coefficient ranges share [`DecompositionStorage::coefficients`].
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(crate) enum CoefficientLayout {
    /// Every component has its own ranges, and all are live together.
    #[default]
    Tile,
    /// Each component's ranges start at zero in one buffer sized to the
    /// largest component. A component's Tier-1, IDWT and store must finish
    /// before the next component is decoded into the buffer.
    Component,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum BuildWorkspace {
    DecodePixels {
        skipped_resolution_levels: u8,
        coefficient_layout: CoefficientLayout,
    },
    CoefficientsOnly,
    /// Plans code blocks without allocating coefficient storage. Sub-band
    /// coefficient ranges are logical only; nothing may index storage by them.
    MetadataOnly,
}

impl BuildWorkspace {
    pub(super) fn active_decomposition_count(self, total_count: usize) -> Option<usize> {
        match self {
            Self::DecodePixels {
                skipped_resolution_levels,
                ..
            } => Some(total_count.saturating_sub(usize::from(skipped_resolution_levels))),
            Self::CoefficientsOnly | Self::MetadataOnly => None,
        }
    }

    pub(crate) fn coefficient_layout(self) -> CoefficientLayout {
        match self {
            Self::DecodePixels {
                coefficient_layout, ..
            } => coefficient_layout,
            Self::CoefficientsOnly | Self::MetadataOnly => CoefficientLayout::Tile,
        }
    }
}

pub(super) struct DecompositionAllocationPlan {
    retained_baseline_bytes: usize,
    pub(super) total_bytes: usize,
    pub(super) coefficients: usize,
    pub(super) tile_decompositions: usize,
    pub(super) decompositions: usize,
    pub(super) sub_bands: usize,
    pub(super) precincts: usize,
    pub(super) code_blocks: usize,
    pub(super) layers: usize,
    /// Layer records of the code-blocks a region decode only parses. They are
    /// never built but are charged like built ones, so the budget rejects a
    /// tile exactly where building every code-block's layers would.
    pub(super) parse_only_layers: usize,
    pub(super) tag_tree_nodes: usize,
    /// Code-blocks and tag-tree nodes of the precincts a region decode
    /// parses in scratch. Like `parse_only_layers`, they are charged but
    /// never built, except for the scratch moved to the built counts.
    pub(super) parse_only_code_blocks: usize,
    pub(super) parse_only_tag_tree_nodes: usize,
    /// The largest precinct parsed in each scratch slot: its code-blocks,
    /// and the nodes of one tag tree, zero when the slot is unused.
    /// LL and HL use slot 0, LH slot 1, and HH slot 2.
    pub(super) scratch_code_blocks: [usize; SCRATCH_PRECINCTS],
    pub(super) scratch_tree_nodes: [usize; SCRATCH_PRECINCTS],
    /// Precincts a region decode never parses. They get no precinct record
    /// or tag-tree roots, but both are charged, so the budget rejects a
    /// tile where storing every precinct would.
    pub(super) unbuilt_precincts: usize,
    exact_integer_decode: bool,
    include_roi_workspace: bool,
    /// Region decode plans coefficients and the whole-plane IDWT workspace
    /// after the ROI plan, so neither is part of this structural plan.
    defer_coefficients: bool,
    workspace: BuildWorkspace,
}

impl DecompositionAllocationPlan {
    fn new(
        retained_baseline_bytes: usize,
        exact_integer_decode: bool,
        include_roi_workspace: bool,
        workspace: BuildWorkspace,
    ) -> Result<Self> {
        if retained_baseline_bytes > DEFAULT_MAX_DECODE_BYTES {
            return Err(ValidationError::ImageTooLarge.into());
        }
        Ok(Self {
            retained_baseline_bytes,
            total_bytes: retained_baseline_bytes,
            coefficients: 0,
            tile_decompositions: 0,
            decompositions: 0,
            sub_bands: 0,
            precincts: 0,
            code_blocks: 0,
            layers: 0,
            parse_only_layers: 0,
            tag_tree_nodes: 0,
            parse_only_code_blocks: 0,
            parse_only_tag_tree_nodes: 0,
            scratch_code_blocks: [0; SCRATCH_PRECINCTS],
            scratch_tree_nodes: [0; SCRATCH_PRECINCTS],
            unbuilt_precincts: 0,
            exact_integer_decode,
            include_roi_workspace,
            defer_coefficients: include_roi_workspace
                && !exact_integer_decode
                && matches!(workspace, BuildWorkspace::DecodePixels { .. }),
            workspace,
        })
    }

    pub(super) fn coefficient_storage_len(&self) -> usize {
        if self.workspace == BuildWorkspace::MetadataOnly || self.defer_coefficients {
            0
        } else {
            self.coefficients
        }
    }

    fn plans_idwt_workspace(&self) -> bool {
        matches!(self.workspace, BuildWorkspace::DecodePixels { .. }) && !self.defer_coefficients
    }

    /// Plan one component's `count` coefficients under the layout.
    fn add_component_coefficients(&mut self, count: usize) -> Result<()> {
        self.coefficients = match self.workspace.coefficient_layout() {
            CoefficientLayout::Tile => self
                .coefficients
                .checked_add(count)
                .ok_or(ValidationError::ImageTooLarge)?,
            CoefficientLayout::Component => self.coefficients.max(count),
        };
        Ok(())
    }

    fn checked_add_count(target_count: &mut usize, added_count: usize) -> Result<()> {
        *target_count = target_count
            .checked_add(added_count)
            .ok_or(ValidationError::ImageTooLarge)?;
        Ok(())
    }

    /// Charge the unbuilt precincts' records and tag-tree roots.
    fn include_unbuilt_precincts(&self, total_bytes: &mut usize) -> Result<()> {
        checked_include_elements::<Precinct>(total_bytes, self.unbuilt_precincts)?;
        let roots = self
            .unbuilt_precincts
            .checked_mul(2)
            .ok_or(ValidationError::ImageTooLarge)?;
        checked_include_elements::<TagNode>(total_bytes, roots)
    }

    /// Reject a logically impossible request before walking every precinct.
    fn validate_minimum_live_workspace(&self, idwt_workspace_elements: usize) -> Result<()> {
        let mut total_bytes = self.retained_baseline_bytes;
        checked_include_elements::<TileDecompositions>(&mut total_bytes, self.tile_decompositions)?;
        checked_include_elements::<Decomposition>(&mut total_bytes, self.decompositions)?;
        checked_include_elements::<SubBand>(&mut total_bytes, self.sub_bands)?;
        checked_include_elements::<Precinct>(&mut total_bytes, self.precincts)?;
        checked_include_elements::<CodeBlock>(&mut total_bytes, self.code_blocks)?;
        checked_include_elements::<Layer>(&mut total_bytes, self.layers)?;
        checked_include_elements::<Layer>(&mut total_bytes, self.parse_only_layers)?;
        checked_include_elements::<TagNode>(&mut total_bytes, self.tag_tree_nodes)?;
        checked_include_elements::<CodeBlock>(&mut total_bytes, self.parse_only_code_blocks)?;
        checked_include_elements::<TagNode>(&mut total_bytes, self.parse_only_tag_tree_nodes)?;
        self.include_unbuilt_precincts(&mut total_bytes)?;
        checked_include_elements::<f32>(&mut total_bytes, self.coefficient_storage_len())?;
        if self.exact_integer_decode {
            checked_include_elements::<i64>(&mut total_bytes, self.coefficient_storage_len())?;
        }

        if self.include_roi_workspace {
            checked_include_elements::<Option<IntRect>>(&mut total_bytes, self.sub_bands)?;
            checked_include_elements::<Option<IntRect>>(&mut total_bytes, self.decompositions)?;
            checked_include_elements::<Option<IntRect>>(
                &mut total_bytes,
                self.tile_decompositions,
            )?;
        }

        if self.plans_idwt_workspace() {
            if self.exact_integer_decode {
                checked_include_elements::<i64>(&mut total_bytes, idwt_workspace_elements)?;
            } else {
                checked_include_elements::<f32>(&mut total_bytes, idwt_workspace_elements)?;
            }
        }
        Ok(())
    }

    fn account_live_workspace(
        &mut self,
        storage: &DecompositionStorage<'_>,
        idwt_workspace_elements: usize,
    ) -> Result<()> {
        let mut total_bytes = self.retained_baseline_bytes;
        checked_include_reusable_elements::<TileDecompositions>(
            &mut total_bytes,
            storage.tile_decompositions.capacity(),
            self.tile_decompositions,
        )?;
        checked_include_reusable_elements::<Decomposition>(
            &mut total_bytes,
            storage.decompositions.capacity(),
            self.decompositions,
        )?;
        checked_include_reusable_elements::<SubBand>(
            &mut total_bytes,
            storage.sub_bands.capacity(),
            self.sub_bands,
        )?;
        checked_include_reusable_elements::<Precinct>(
            &mut total_bytes,
            storage.precincts.capacity(),
            self.precincts,
        )?;
        checked_include_reusable_elements::<CodeBlock>(
            &mut total_bytes,
            storage.code_blocks.capacity(),
            self.code_blocks,
        )?;
        checked_include_reusable_elements::<Layer>(
            &mut total_bytes,
            storage.layers.capacity(),
            self.layers,
        )?;
        checked_include_elements::<Layer>(&mut total_bytes, self.parse_only_layers)?;
        checked_include_reusable_elements::<TagNode>(
            &mut total_bytes,
            storage.tag_tree_nodes.capacity(),
            self.tag_tree_nodes,
        )?;
        checked_include_elements::<CodeBlock>(&mut total_bytes, self.parse_only_code_blocks)?;
        checked_include_elements::<TagNode>(&mut total_bytes, self.parse_only_tag_tree_nodes)?;
        self.include_unbuilt_precincts(&mut total_bytes)?;
        checked_include_reusable_elements::<f32>(
            &mut total_bytes,
            storage.coefficients.capacity(),
            self.coefficient_storage_len(),
        )?;
        checked_include_reusable_elements::<i64>(
            &mut total_bytes,
            storage.coefficients_i64.capacity(),
            usize::from(self.exact_integer_decode) * self.coefficient_storage_len(),
        )?;

        if self.include_roi_workspace {
            if let Some(roi_plan) = storage.roi_plan.as_ref() {
                total_bytes = total_bytes
                    .checked_add(roi_plan.owned_bytes()?)
                    .filter(|bytes| *bytes <= DEFAULT_MAX_DECODE_BYTES)
                    .ok_or(ValidationError::ImageTooLarge)?;
            } else {
                checked_include_elements::<Option<IntRect>>(&mut total_bytes, self.sub_bands)?;
                checked_include_elements::<Option<IntRect>>(&mut total_bytes, self.decompositions)?;
                checked_include_elements::<Option<IntRect>>(
                    &mut total_bytes,
                    self.tile_decompositions,
                )?;
            }
        }

        if self.plans_idwt_workspace() {
            if self.exact_integer_decode {
                checked_include_elements::<i64>(&mut total_bytes, idwt_workspace_elements)?;
            } else {
                checked_include_elements::<f32>(&mut total_bytes, idwt_workspace_elements)?;
            }
        }

        // The packet parser owns segment growth under the remaining budget.
        // Validate retained capacity once, but keep it outside the structural
        // total so downstream phases add this owner exactly once.
        let mut total_with_segments = total_bytes;
        checked_include_elements::<Segment<'_>>(
            &mut total_with_segments,
            storage.segments.capacity(),
        )?;
        self.total_bytes = total_bytes;
        Ok(())
    }
}

pub(super) fn prepare_decomposition_storage(
    tile: &Tile<'_>,
    storage: &mut DecompositionStorage<'_>,
    retained_baseline_bytes: usize,
    include_roi_workspace: bool,
    workspace: BuildWorkspace,
    selection: Option<&super::PrecinctSelection>,
) -> Result<DecompositionAllocationPlan> {
    let (mut plan, idwt_workspace_elements) = plan::build_allocation_plan(
        tile,
        retained_baseline_bytes,
        storage.exact_integer_decode,
        include_roi_workspace,
        workspace,
        selection,
        // A selection always comes with the ROI plan it was made from.
        selection.and(storage.roi_plan.as_ref()),
    )?;

    // Capacity retention must never make a smaller valid tile fail in a later
    // phase. Keep exact/current-size buffers, discard stale excess, and start
    // packet metadata empty because its target is not known until parsing.
    reuse::discard_stale_capacity(storage, &plan);
    plan.account_live_workspace(storage, idwt_workspace_elements)?;
    reuse::reserve_decomposition_storage(&plan, storage, retained_baseline_bytes)?;
    plan.account_live_workspace(storage, idwt_workspace_elements)?;
    storage.deferred_idwt_workspace_elements =
        plan.defer_coefficients.then_some(idwt_workspace_elements);
    Ok(plan)
}

/// Allocate coefficient storage that region decode deferred until its ROI
/// plan was built.
///
/// With a plan, each sub-band stores only the union of the code-blocks the
/// plan requires; every other code-block would be left zero anyway, and
/// [`super::SubBand::storage_rect`] readers treat samples outside the stored
/// rectangle as zero. With `clip_to_idwt_windows`, that union is further
/// clipped to the window the IDWT reads, so code-blocks on the window's edge
/// keep only the samples inside it. A code-block decoder hook writes whole
/// code-blocks, so hooked decodes keep the whole union. The windowed IDWT's
/// peak buffers are charged in place of the whole-plane workspace. Without a
/// plan, the whole tile is stored and the whole-plane IDWT workspace is
/// charged as an ordinary decode would.
///
/// Like the whole-plane workspace, the windowed IDWT charge stays in
/// `structural_workspace_bytes` for the rest of the tile, so packet parsing
/// and Tier-1 budget against it before any IDWT buffer exists.
pub(crate) fn allocate_deferred_coefficients(
    storage: &mut DecompositionStorage<'_>,
    skipped_resolution_levels: u8,
    clip_to_idwt_windows: bool,
) -> Result<()> {
    let Some(idwt_workspace_elements) = storage.deferred_idwt_workspace_elements.take() else {
        return Ok(());
    };
    let mut live_bytes = storage.structural_workspace_bytes;
    let coefficient_count = if storage.roi_plan.is_some() {
        let idwt_elements =
            crate::j2c::idwt::roi_idwt_peak_elements(storage, skipped_resolution_levels)?;
        storage.roi_idwt_charge_elements = idwt_elements;
        checked_include_elements::<f32>(&mut live_bytes, idwt_elements)?;
        lay_out_region_coefficients(storage, clip_to_idwt_windows)?
    } else {
        checked_include_elements::<f32>(&mut live_bytes, idwt_workspace_elements)?;
        storage
            .sub_bands
            .iter()
            .map(|sub_band| sub_band.coefficients.end)
            .max()
            .unwrap_or(0)
    };
    let mut budget = reuse::ReallocationBudget::from_live_bytes(live_bytes)?;
    budget.reserve(&mut storage.coefficients, coefficient_count)?;
    crate::try_resize_decode_elements(&mut storage.coefficients, coefficient_count, 0.0)?;
    storage.structural_workspace_bytes = budget.live_bytes();
    Ok(())
}

/// Give each sub-band of a region decode the coefficient range for its
/// stored rectangle: the required code-blocks, clipped to the IDWT windows
/// unless `clip_to_idwt_windows` is false. With banded Tier-1, a range holds
/// one band of the rectangle's rows. Returns the coefficients the ranges
/// need.
fn lay_out_region_coefficients(
    storage: &mut DecompositionStorage<'_>,
    clip_to_idwt_windows: bool,
) -> Result<usize> {
    let Some(roi_plan) = storage.roi_plan.as_ref() else {
        return Err(DecodingError::CodeBlockDecodeFailure.into());
    };
    let mut coefficient_count = 0usize;
    // Build stores each component's sub-bands together, starting with its
    // LL band, so a component starts where its LL band does.
    let per_component = storage.coefficient_layout == CoefficientLayout::Component;
    let mut component_starts = storage
        .tile_decompositions
        .iter()
        .map(|component| component.first_ll_sub_band)
        .peekable();
    let mut next_offset = 0usize;
    for (sub_band_idx, sub_band) in storage.sub_bands.iter_mut().enumerate() {
        if component_starts.next_if_eq(&sub_band_idx).is_some() && per_component {
            next_offset = 0;
        }
        let mut stored: Option<IntRect> = None;
        for precinct in &storage.precincts[sub_band.precincts.clone()] {
            for code_block in &storage.code_blocks[precinct.code_blocks.clone()] {
                if roi_plan.code_block_required(sub_band_idx, code_block.rect) {
                    stored =
                        Some(stored.map_or(code_block.rect, |rect| rect.union(code_block.rect)));
                }
            }
        }
        let stored = stored
            .and_then(|stored| {
                if !clip_to_idwt_windows {
                    return Some(stored);
                }
                roi_plan
                    .sub_band_window(sub_band_idx)
                    .filter(|window| window.intersects(stored))
                    .map(|window| window.intersect(stored))
            })
            .unwrap_or_else(|| IntRect::from_xywh(sub_band.rect.x0, sub_band.rect.y0, 0, 0));
        // Banded Tier-1 holds one band of code-block rows at a time: as
        // many rows as any band shares with the stored rectangle.
        let rows = if storage.banded_tier1 {
            storage.precincts[sub_band.precincts.clone()]
                .iter()
                .flat_map(|precinct| &storage.code_blocks[precinct.code_blocks.clone()])
                .map(|code_block| {
                    code_block
                        .rect
                        .y1
                        .min(stored.y1)
                        .saturating_sub(code_block.rect.y0.max(stored.y0))
                })
                .max()
                .unwrap_or(0)
        } else {
            stored.height()
        };
        let count = (stored.width() as usize)
            .checked_mul(rows as usize)
            .ok_or(ValidationError::ImageTooLarge)?;
        let end = next_offset
            .checked_add(count)
            .ok_or(ValidationError::ImageTooLarge)?;
        sub_band.coefficients = next_offset..end;
        sub_band.storage_rect = stored;
        next_offset = end;
        coefficient_count = coefficient_count.max(end);
    }
    Ok(coefficient_count)
}

/// After a banded Tier-1 failed, store every sub-band's whole rectangle
/// again, within the cap, so Tier-1 can rerun the old way.
pub(crate) fn lay_out_whole_region_coefficients(
    storage: &mut DecompositionStorage<'_>,
) -> Result<()> {
    storage.banded_tier1 = false;
    let coefficient_count = lay_out_region_coefficients(storage, true)?;
    let mut budget =
        reuse::ReallocationBudget::from_live_bytes(storage.structural_workspace_bytes)?;
    budget.reserve(&mut storage.coefficients, coefficient_count)?;
    crate::try_resize_decode_elements(&mut storage.coefficients, coefficient_count, 0.0)?;
    storage.structural_workspace_bytes = budget.live_bytes();
    Ok(())
}

pub(crate) fn release_unused_roi_workspace(
    storage: &mut DecompositionStorage<'_>,
    component_count: usize,
) -> Result<()> {
    let roi_bytes = roi_workspace_bytes(
        storage.sub_bands.len(),
        storage.decompositions.len(),
        component_count,
    )?;
    release_roi_workspace_bytes(&mut storage.structural_workspace_bytes, roi_bytes)
}

fn roi_workspace_bytes(
    sub_band_count: usize,
    decomposition_count: usize,
    component_count: usize,
) -> Result<usize> {
    let mut roi_bytes = 0usize;
    checked_include_elements::<Option<IntRect>>(&mut roi_bytes, sub_band_count)?;
    checked_include_elements::<Option<IntRect>>(&mut roi_bytes, decomposition_count)?;
    checked_include_elements::<Option<IntRect>>(&mut roi_bytes, component_count)?;
    Ok(roi_bytes)
}

fn release_roi_workspace_bytes(
    structural_workspace_bytes: &mut usize,
    roi_bytes: usize,
) -> Result<()> {
    *structural_workspace_bytes = structural_workspace_bytes
        .checked_sub(roi_bytes)
        .ok_or(ValidationError::ImageTooLarge)?;
    Ok(())
}

pub(super) fn checked_include_elements<T>(total_bytes: &mut usize, count: usize) -> Result<()> {
    let bytes = size_of::<T>()
        .checked_mul(count)
        .ok_or(ValidationError::ImageTooLarge)?;
    *total_bytes = total_bytes
        .checked_add(bytes)
        .ok_or(ValidationError::ImageTooLarge)?;
    if *total_bytes > DEFAULT_MAX_DECODE_BYTES {
        return Err(ValidationError::ImageTooLarge.into());
    }
    Ok(())
}

fn checked_include_reusable_elements<T>(
    total_bytes: &mut usize,
    retained_capacity: usize,
    target_count: usize,
) -> Result<()> {
    checked_include_elements::<T>(total_bytes, retained_capacity.max(target_count))
}

#[cfg(test)]
mod tests;