j2k-native 0.13.0

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

//! Move-only precinct partitioning for prepared Tier-1 owners.

use alloc::vec::Vec;

use super::super::{
    NativeEncodePipelineError, NativeEncodePipelineResult, NativeEncodeSession,
    PreparedEncodeSubband, PreparedResolutionPacket,
};

mod geometry;
use geometry::{
    component_bounds, component_split_packet_count, precinct_exponents_for_resolution,
    resolution_precinct_grid,
};
mod distribution;
use distribution::distribute_owned_subband;
mod ownership;
use ownership::{try_destination_vec, try_push_planned, PrecinctSplitAccounting};

pub(in crate::j2c::encode) fn split_component_resolution_packets_by_precinct_for_session(
    component_resolution_packets: Vec<Vec<PreparedResolutionPacket>>,
    tile: j2k_types::J2kRect,
    component_sampling: &[(u8, u8)],
    num_decomposition_levels: u8,
    precinct_exponents: &[(u8, u8)],
    session: &NativeEncodeSession<'_>,
    retained_phase_bytes: usize,
) -> NativeEncodePipelineResult<Vec<Vec<PreparedResolutionPacket>>> {
    let (packets, _peak_phase_bytes) = split_component_resolution_packets_by_precinct_accounted(
        component_resolution_packets,
        tile,
        component_sampling,
        num_decomposition_levels,
        precinct_exponents,
        session,
        retained_phase_bytes,
    )?;
    Ok(packets)
}

fn split_component_resolution_packets_by_precinct_accounted(
    component_resolution_packets: Vec<Vec<PreparedResolutionPacket>>,
    tile: j2k_types::J2kRect,
    component_sampling: &[(u8, u8)],
    num_decomposition_levels: u8,
    precinct_exponents: &[(u8, u8)],
    session: &NativeEncodeSession<'_>,
    retained_phase_bytes: usize,
) -> NativeEncodePipelineResult<(Vec<Vec<PreparedResolutionPacket>>, usize)> {
    let source_capacity = component_resolution_packets.capacity();
    let mut accounting = PrecinctSplitAccounting::try_from_source(
        session,
        retained_phase_bytes,
        &component_resolution_packets,
        source_capacity,
    )?;
    if precinct_exponents.is_empty() {
        // COD's omitted precinct sizes mean 2^15 at every resolution, not an
        // unbounded precinct. Keep the common one-packet route allocation-free.
        let mut needs_split = false;
        for (component, packets) in component_resolution_packets.iter().enumerate() {
            let bounds = component_bounds(tile, component_sampling, component)?;
            for packet in packets {
                let grid = resolution_precinct_grid(
                    bounds,
                    num_decomposition_levels,
                    packet.resolution,
                    15,
                    15,
                )?;
                needs_split |= grid.columns != 1 || grid.rows != 1;
            }
        }
        if !needs_split {
            return Ok((component_resolution_packets, accounting.peak_phase_bytes()));
        }
    }

    let component_count = component_resolution_packets.len();
    let mut split_components = try_destination_vec(
        component_count,
        &mut accounting,
        "split component packet owners",
    )?;
    for (component, component_packets) in component_resolution_packets.into_iter().enumerate() {
        // Sampled components have smaller resolutions, so their precinct
        // grids come from their own tile-component extent.
        let bounds = component_bounds(tile, component_sampling, component)?;
        let component_capacity = component_packets.capacity();
        let split_packet_count = component_split_packet_count(
            &component_packets,
            bounds,
            num_decomposition_levels,
            precinct_exponents,
        )?;
        let mut split_packets = try_destination_vec(
            split_packet_count,
            &mut accounting,
            "split resolution packet owners",
        )?;
        for packet in component_packets {
            split_prepared_resolution_packet_by_precinct(
                packet,
                bounds,
                num_decomposition_levels,
                precinct_exponents,
                &mut split_packets,
                &mut accounting,
            )?;
        }
        accounting.release_source_capacity::<PreparedResolutionPacket>(
            component_capacity,
            "source resolution packet owners",
        )?;
        try_push_planned(&mut split_components, split_packets)?;
    }
    accounting.release_source_capacity::<Vec<PreparedResolutionPacket>>(
        source_capacity,
        "source component packet owners",
    )?;
    let peak_phase_bytes = accounting.finish(&split_components, split_components.capacity())?;
    Ok((split_components, peak_phase_bytes))
}

fn split_prepared_resolution_packet_by_precinct(
    packet: PreparedResolutionPacket,
    bounds: j2k_types::J2kRect,
    num_decomposition_levels: u8,
    precinct_exponents: &[(u8, u8)],
    split_packets: &mut Vec<PreparedResolutionPacket>,
    accounting: &mut PrecinctSplitAccounting<'_, '_>,
) -> NativeEncodePipelineResult<()> {
    let PreparedResolutionPacket {
        component,
        resolution,
        precinct: _,
        subbands,
    } = packet;
    let (horizontal_exponent, vertical_exponent) =
        precinct_exponents_for_resolution(precinct_exponents, resolution)?;
    let grid = resolution_precinct_grid(
        bounds,
        num_decomposition_levels,
        resolution,
        horizontal_exponent,
        vertical_exponent,
    )?;
    let packet_start = split_packets.len();
    let subband_count = subbands.len();
    for precinct_row in 0..grid.rows {
        for precinct_column in 0..grid.columns {
            let precinct = u64::from(precinct_row)
                .checked_mul(u64::from(grid.columns))
                .and_then(|value| value.checked_add(u64::from(precinct_column)))
                .ok_or_else(|| {
                    NativeEncodePipelineError::arithmetic_overflow("precinct index overflow")
                })?;
            let split_subbands =
                try_destination_vec(subband_count, accounting, "split prepared subband owners")?;
            try_push_planned(
                split_packets,
                PreparedResolutionPacket {
                    component,
                    resolution,
                    precinct,
                    subbands: split_subbands,
                },
            )?;
        }
    }

    let source_subband_capacity = subbands.capacity();
    let packet_end = split_packets.len();
    let packet_slice = split_packets
        .get_mut(packet_start..packet_end)
        .ok_or_else(|| {
            NativeEncodePipelineError::internal_invariant("precinct packet output range mismatch")
        })?;
    for subband in subbands {
        distribute_owned_subband(
            subband,
            resolution,
            horizontal_exponent,
            vertical_exponent,
            grid,
            packet_slice,
            accounting,
        )?;
    }
    accounting.release_source_capacity::<PreparedEncodeSubband>(
        source_subband_capacity,
        "source prepared subband owners",
    )?;
    Ok(())
}

#[cfg(test)]
mod tests;