j2k-native 0.7.0

//! The irreversible multi-component transformation, as specified in
//! Annex G.2 and G.3.

use super::codestream::{ComponentInfo, Header, WaveletTransform};
use super::decode::TileDecodeContext;
use crate::error::{bail, err, ColorError, Result};
use crate::math::{dispatch, f32x8, floor_f32, Level, Simd};
use crate::{HtCodeBlockDecoder, J2kInverseMctJob, J2kWaveletTransform};
use j2k_codec_math::mct;

/// Apply the inverse multi-component transform, as specified in G.2 and G.3.
pub(crate) fn apply_inverse(
    tile_ctx: &mut TileDecodeContext,
    component_infos: &[super::codestream::ComponentInfo],
    header: &Header<'_>,
    backend: &mut Option<&mut dyn HtCodeBlockDecoder>,
) -> Result<()> {
    if tile_ctx.channel_data.len() < 3 {
        return if header.strict {
            err!(ColorError::Mct)
        } else {
            Ok(())
        };
    }

    let (s, _) = tile_ctx.channel_data.split_at_mut(3);
    let [s0, s1, s2] = s else { unreachable!() };

    let transform = component_infos[0].wavelet_transform();

    if transform != component_infos[1].wavelet_transform()
        || component_infos[1].wavelet_transform() != component_infos[2].wavelet_transform()
    {
        bail!(ColorError::Mct);
    }

    if s0.container.len() != s1.container.len() || s1.container.len() != s2.container.len() {
        bail!(ColorError::Mct);
    }

    if s0.integer_container.is_some()
        || s1.integer_container.is_some()
        || s2.integer_container.is_some()
    {
        return apply_inverse_i64(transform, s0, s1, s2);
    }

    let handled = if let Some(backend) = backend.as_deref_mut() {
        backend.decode_inverse_mct(J2kInverseMctJob {
            transform: J2kWaveletTransform::from(transform),
            plane0: &mut s0.container,
            plane1: &mut s1.container,
            plane2: &mut s2.container,
            addend0: unsigned_level_shift(&component_infos[0]),
            addend1: unsigned_level_shift(&component_infos[1]),
            addend2: unsigned_level_shift(&component_infos[2]),
        })?
    } else {
        false
    };

    if !handled {
        apply_inner(
            transform,
            &mut s0.container,
            &mut s1.container,
            &mut s2.container,
        );
    }

    Ok(())
}

fn apply_inverse_i64(
    transform: WaveletTransform,
    s0: &mut super::ComponentData,
    s1: &mut super::ComponentData,
    s2: &mut super::ComponentData,
) -> Result<()> {
    if transform != WaveletTransform::Reversible53 {
        bail!(ColorError::Mct);
    }

    let (Some(y0), Some(y1), Some(y2)) = (
        s0.integer_container.as_mut(),
        s1.integer_container.as_mut(),
        s2.integer_container.as_mut(),
    ) else {
        bail!(ColorError::Mct);
    };
    if y0.len() != y1.len() || y1.len() != y2.len() {
        bail!(ColorError::Mct);
    }

    for ((y0, y1), y2) in y0.iter_mut().zip(y1.iter_mut()).zip(y2.iter_mut()) {
        let src0 = *y0;
        let src1 = *y1;
        let src2 = *y2;
        let green = src0 - floor_div_i64(src2 + src1, 4);
        *y0 = src2 + green;
        *y1 = green;
        *y2 = src1 + green;
    }

    Ok(())
}

#[expect(
    clippy::cast_precision_loss,
    reason = "the codec float domain intentionally receives bounded integer samples or metadata at this rounding boundary"
)]
fn unsigned_level_shift(component_info: &ComponentInfo) -> f32 {
    if component_info.size_info.signed {
        0.0
    } else {
        (1_u32 << (component_info.size_info.precision - 1)) as f32
    }
}

#[expect(
    clippy::inline_always,
    reason = "this scalar primitive is intentionally inlined into the reversible color-transform hot loop"
)]
#[inline(always)]
fn floor_div_i64(numerator: i64, denominator: i64) -> i64 {
    debug_assert!(denominator > 0);
    let quotient = numerator / denominator;
    let remainder = numerator % denominator;
    if remainder != 0 && remainder < 0 {
        quotient - 1
    } else {
        quotient
    }
}

fn apply_inner(transform: WaveletTransform, s0: &mut [f32], s1: &mut [f32], s2: &mut [f32]) {
    dispatch!(Level::new(), simd => apply_inner_impl(simd, transform, s0, s1, s2));
}

#[expect(
    clippy::inline_always,
    reason = "the SIMD implementation is intentionally specialized and inlined at the architecture dispatch boundary"
)]
#[inline(always)]
fn apply_inner_impl<S: Simd>(
    simd: S,
    transform: WaveletTransform,
    s0: &mut [f32],
    s1: &mut [f32],
    s2: &mut [f32],
) {
    match transform {
        // Irreversible MCT, specified in G.3.
        WaveletTransform::Irreversible97 => {
            let mut s0_chunks = s0.chunks_exact_mut(8);
            let mut s1_chunks = s1.chunks_exact_mut(8);
            let mut s2_chunks = s2.chunks_exact_mut(8);
            for ((y0, y1), y2) in s0_chunks
                .by_ref()
                .zip(s1_chunks.by_ref())
                .zip(s2_chunks.by_ref())
            {
                let y_0 = f32x8::from_slice(simd, y0);
                let y_1 = f32x8::from_slice(simd, y1);
                let y_2 = f32x8::from_slice(simd, y2);

                let i0 = y_2.mul_add(f32x8::splat(simd, mct::ICT_INV_R_CR), y_0);
                let i1 = y_2.mul_add(
                    f32x8::splat(simd, mct::ICT_INV_G_CR),
                    y_1.mul_add(f32x8::splat(simd, mct::ICT_INV_G_CB), y_0),
                );
                let i2 = y_1.mul_add(f32x8::splat(simd, mct::ICT_INV_B_CB), y_0);

                i0.store(y0);
                i1.store(y1);
                i2.store(y2);
            }
            for ((y0, y1), y2) in s0_chunks
                .into_remainder()
                .iter_mut()
                .zip(s1_chunks.into_remainder().iter_mut())
                .zip(s2_chunks.into_remainder().iter_mut())
            {
                let src0 = *y0;
                let src1 = *y1;
                let src2 = *y2;
                *y0 = src0 + mct::ICT_INV_R_CR * src2;
                *y1 = src0 + mct::ICT_INV_G_CB * src1 + mct::ICT_INV_G_CR * src2;
                *y2 = src0 + mct::ICT_INV_B_CB * src1;
            }
        }
        // Reversible MCT, specified in G.2.
        WaveletTransform::Reversible53 => {
            let mut s0_chunks = s0.chunks_exact_mut(8);
            let mut s1_chunks = s1.chunks_exact_mut(8);
            let mut s2_chunks = s2.chunks_exact_mut(8);
            for ((y0, y1), y2) in s0_chunks
                .by_ref()
                .zip(s1_chunks.by_ref())
                .zip(s2_chunks.by_ref())
            {
                let y_0 = f32x8::from_slice(simd, y0);
                let y_1 = f32x8::from_slice(simd, y1);
                let y_2 = f32x8::from_slice(simd, y2);

                let i1 = y_0 - ((y_2 + y_1) * mct::RCT_QUARTER).floor();
                let i0 = y_2 + i1;
                let i2 = y_1 + i1;

                i0.store(y0);
                i1.store(y1);
                i2.store(y2);
            }
            for ((y0, y1), y2) in s0_chunks
                .into_remainder()
                .iter_mut()
                .zip(s1_chunks.into_remainder().iter_mut())
                .zip(s2_chunks.into_remainder().iter_mut())
            {
                let src0 = *y0;
                let src1 = *y1;
                let src2 = *y2;
                let i1 = src0 - floor_f32((src2 + src1) * mct::RCT_QUARTER);
                *y0 = src2 + i1;
                *y1 = i1;
                *y2 = src1 + i1;
            }
        }
    }
}