j2k-native 0.13.1

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

//! HT cleanup-magnitude aggregation before prepared coefficients are consumed.

use super::{
    ht_block_encode, BlockCodingMode, NativeEncodePipelineError, NativeEncodePipelineResult,
    PreparedCodeBlockCoefficients, PreparedEncodeCodeBlock, PreparedEncodeSubband,
    PreparedResolutionPacket,
};
use crate::j2c::capabilities::required_magnitude_bound;

pub(super) fn decomposition_level_for_resolution(
    resolution: impl TryInto<u8>,
    num_decomposition_levels: u8,
) -> NativeEncodePipelineResult<u8> {
    let resolution = resolution.try_into().map_err(|_| {
        NativeEncodePipelineError::internal_invariant(
            "resolution index exceeds the HT magnitude-bound domain",
        )
    })?;
    if resolution == 0 {
        Ok(num_decomposition_levels)
    } else {
        num_decomposition_levels
            .checked_sub(resolution - 1)
            .ok_or_else(|| {
                NativeEncodePipelineError::internal_invariant(
                    "resolution exceeds the decomposition count",
                )
            })
    }
}

/// `parallel` lets each subband's code blocks be scanned on the Rayon pool.
pub(super) fn required_ht_magnitude_bound<'a>(
    packets: impl IntoIterator<Item = &'a PreparedResolutionPacket>,
    num_decomposition_levels: u8,
    reversible: bool,
    parallel: bool,
) -> NativeEncodePipelineResult<Option<u8>> {
    let mut required = None::<u8>;
    for packet in packets {
        let decomposition_level =
            decomposition_level_for_resolution(packet.resolution, num_decomposition_levels)?;
        for subband in &packet.subbands {
            if subband.block_coding_mode != BlockCodingMode::HighThroughput {
                continue;
            }
            let maximum = maximum_cleanup_magnitude(subband, parallel);
            let subband_required =
                required_magnitude_bound(maximum, reversible, decomposition_level);
            required = Some(required.map_or(subband_required, |bound| bound.max(subband_required)));
        }
    }
    Ok(required)
}

pub(super) fn cleanup_magnitude_upper_bound(
    total_bitplanes: u8,
    num_zero_bitplanes: u8,
    num_coding_passes: u8,
) -> u64 {
    let cleanup_bitplanes = total_bitplanes
        .saturating_sub(num_zero_bitplanes)
        .saturating_sub(u8::from(num_coding_passes >= 2));
    if u32::from(cleanup_bitplanes) >= u64::BITS {
        u64::MAX
    } else if cleanup_bitplanes == 0 {
        0
    } else {
        (1_u64 << cleanup_bitplanes) - 1
    }
}

/// A code block's largest magnitude, and whether it holds coefficients.
fn code_block_maximum(block: &PreparedEncodeCodeBlock) -> (u64, bool) {
    match &block.coefficients {
        // A u32 maximum vectorizes; widening first does not.
        PreparedCodeBlockCoefficients::I32(values) => (
            u64::from(
                values
                    .iter()
                    .map(|value| value.unsigned_abs())
                    .max()
                    .unwrap_or(0),
            ),
            true,
        ),
        PreparedCodeBlockCoefficients::I64(values) => (
            values
                .iter()
                .map(|value| value.unsigned_abs())
                .max()
                .unwrap_or(0),
            true,
        ),
        PreparedCodeBlockCoefficients::Empty => (0, false),
    }
}

fn maximum_cleanup_magnitude(subband: &PreparedEncodeSubband, parallel: bool) -> u64 {
    let refinement_shift = u8::from(
        ht_block_encode::effective_coding_passes(
            subband.total_bitplanes,
            subband.ht_target_coding_passes,
        ) >= 2,
    );
    let merge =
        |(a, a_exact): (u64, bool), (b, b_exact): (u64, bool)| (a.max(b), a_exact && b_exact);
    let blocks = &subband.code_blocks;
    #[cfg(feature = "parallel")]
    let (mut maximum, all_blocks_exact) =
        if parallel && blocks.len() > 1 && rayon::current_num_threads() > 1 {
            use rayon::prelude::*;
            let maxima = blocks.par_iter().with_min_len(8).map(code_block_maximum);
            maxima.reduce(|| (0, true), merge)
        } else {
            blocks.iter().map(code_block_maximum).fold((0, true), merge)
        };
    #[cfg(not(feature = "parallel"))]
    let (mut maximum, all_blocks_exact) = {
        let _ = parallel;
        blocks.iter().map(code_block_maximum).fold((0, true), merge)
    };
    let exact = all_blocks_exact && subband.preencoded_ht_code_blocks.is_none();
    maximum >>= u32::from(refinement_shift);
    if let Some(exact) = subband.preencoded_ht_maximum_cleanup_magnitude {
        return maximum.max(exact);
    }
    if let Some(encoded_blocks) = &subband.preencoded_ht_code_blocks {
        let conservative = encoded_blocks
            .iter()
            .map(|block| {
                cleanup_magnitude_upper_bound(
                    subband.total_bitplanes,
                    block.num_zero_bitplanes,
                    block.num_coding_passes,
                )
            })
            .max()
            .unwrap_or(0);
        return maximum.max(conservative);
    }
    if exact {
        return maximum;
    }

    let cleanup_bitplanes = subband.total_bitplanes.saturating_sub(refinement_shift);
    let conservative = if u32::from(cleanup_bitplanes) >= u64::BITS {
        u64::MAX
    } else if cleanup_bitplanes == 0 {
        0
    } else {
        (1_u64 << cleanup_bitplanes) - 1
    };
    maximum.max(conservative)
}