j2k-native 0.13.0

use alloc::vec::Vec;

pub(crate) const SIMD_WIDTH: usize = 8;

/// Number of bits required to represent an unsigned 32-bit magnitude.
///
/// The count is kept in its natural `0..=32` range instead of being
/// calculated in a wider type and narrowed at each call site.
pub(crate) const fn bit_width_u32(mut value: u32) -> u8 {
    let mut bits = 0_u8;
    while value != 0 {
        bits += 1;
        value >>= 1;
    }
    bits
}

/// Number of bits required to represent an unsigned 64-bit magnitude.
pub(crate) const fn bit_width_u64(mut value: u64) -> u8 {
    let mut bits = 0_u8;
    while value != 0 {
        bits += 1;
        value >>= 1;
    }
    bits
}

/// Smallest exponent whose power of two is at least `value`.
pub(crate) const fn ceil_log2_u32(value: u32) -> u8 {
    if value <= 1 {
        0
    } else {
        bit_width_u32(value - 1)
    }
}

#[cfg(feature = "simd")]
mod inner {
    use super::SIMD_WIDTH;
    use core::ops::{Add, AddAssign, DivAssign, Mul, MulAssign, Sub, SubAssign};
    use fearless_simd::{SimdBase, SimdFloat};

    pub(crate) use fearless_simd::{dispatch, Level, Simd};

    #[derive(Copy, Clone)]
    #[repr(C, align(32))]
    pub(crate) struct f32x8<S: Simd> {
        inner: fearless_simd::f32x8<S>,
    }

    impl<S: Simd> f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn from_slice(simd: S, slice: &[f32]) -> Self {
            Self {
                inner: fearless_simd::f32x8::from_slice(simd, slice),
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn splat(simd: S, value: f32) -> Self {
            Self {
                inner: fearless_simd::f32x8::splat(simd, value),
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn mul_add(self, mul: Self, addend: Self) -> Self {
            Self {
                // Preserve the scalar path's single-rounding contract without hardware FMA.
                inner: self.inner.mul_add_precise(mul.inner, addend.inner),
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn floor(self) -> Self {
            Self {
                inner: self.inner.floor(),
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn store(self, slice: &mut [f32]) {
            self.inner.store_slice(&mut slice[..SIMD_WIDTH]);
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn zip_low(self, other: Self) -> Self {
            Self {
                inner: self.inner.zip_low(other.inner),
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn zip_high(self, other: Self) -> Self {
            Self {
                inner: self.inner.zip_high(other.inner),
            }
        }

        /// Lanes `SHIFT..` of `self` followed by the first `SHIFT` lanes of `next`.
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn slide<const SHIFT: usize>(self, next: Self) -> Self {
            Self {
                inner: self.inner.slide::<SHIFT>(next.inner),
            }
        }

        /// The even lanes of `self` then of `next`, and their odd lanes.
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn deinterleave(self, next: Self) -> (Self, Self) {
            let (even, odd) = self.inner.deinterleave(next.inner);
            (Self { inner: even }, Self { inner: odd })
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn min(self, other: Self) -> Self {
            Self {
                inner: self.inner.min(other.inner),
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn max(self, other: Self) -> Self {
            Self {
                inner: self.inner.max(other.inner),
            }
        }
    }

    impl<S: Simd> Add for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn add(self, rhs: Self) -> Self {
            Self {
                inner: self.inner + rhs.inner,
            }
        }
    }

    impl<S: Simd> Sub for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn sub(self, rhs: Self) -> Self {
            Self {
                inner: self.inner - rhs.inner,
            }
        }
    }

    impl<S: Simd> Mul for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn mul(self, rhs: Self) -> Self {
            Self {
                inner: self.inner * rhs.inner,
            }
        }
    }

    impl<S: Simd> Add<f32> for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn add(self, rhs: f32) -> Self {
            Self {
                inner: self.inner + rhs,
            }
        }
    }

    impl<S: Simd> Mul<f32> for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn mul(self, rhs: f32) -> Self {
            Self {
                inner: self.inner * rhs,
            }
        }
    }

    impl<S: Simd> AddAssign for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn add_assign(&mut self, rhs: Self) {
            self.inner = self.inner + rhs.inner;
        }
    }

    impl<S: Simd> SubAssign for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn sub_assign(&mut self, rhs: Self) {
            self.inner = self.inner - rhs.inner;
        }
    }

    impl<S: Simd> MulAssign<f32> for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn mul_assign(&mut self, rhs: f32) {
            self.inner = self.inner * rhs;
        }
    }

    impl<S: Simd> DivAssign<f32> for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn div_assign(&mut self, rhs: f32) {
            self.inner = self.inner / rhs;
        }
    }
}

#[cfg(not(feature = "simd"))]
mod inner {
    use super::SIMD_WIDTH;
    use core::marker::PhantomData;
    use core::ops::{Add, AddAssign, DivAssign, Mul, MulAssign, Sub, SubAssign};

    /// Mirrors `fearless_simd::Simd`, whose tokens are `Send + Sync` so
    /// transform kernels can run on worker threads.
    pub(crate) trait Simd: Copy + Clone + Send + Sync {}

    #[derive(Copy, Clone)]
    pub(crate) struct ScalarSimd;
    impl Simd for ScalarSimd {}

    pub(crate) struct Level;
    impl Level {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn new() -> Self {
            Level
        }
    }

    #[derive(Copy, Clone)]
    #[repr(C, align(32))]
    pub(crate) struct f32x8<S: Simd> {
        val: [f32; SIMD_WIDTH],
        _marker: PhantomData<S>,
    }

    impl<S: Simd> f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn from_slice(_simd: S, slice: &[f32]) -> Self {
            let mut val = [0.0f32; SIMD_WIDTH];
            val.copy_from_slice(&slice[..SIMD_WIDTH]);
            Self {
                val,
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn splat(_simd: S, value: f32) -> Self {
            Self {
                val: [value; SIMD_WIDTH],
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        pub(crate) fn mul_add(self, mul: Self, addend: Self) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = super::mul_add(self.val[i], mul.val[i], addend.val[i]);
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        pub(crate) fn floor(self) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = super::floor_f32(self.val[i]);
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn store(self, slice: &mut [f32]) {
            slice[..SIMD_WIDTH].copy_from_slice(&self.val);
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn zip_low(self, other: Self) -> Self {
            Self {
                val: [
                    self.val[0],
                    other.val[0],
                    self.val[1],
                    other.val[1],
                    self.val[2],
                    other.val[2],
                    self.val[3],
                    other.val[3],
                ],
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn zip_high(self, other: Self) -> Self {
            Self {
                val: [
                    self.val[4],
                    other.val[4],
                    self.val[5],
                    other.val[5],
                    self.val[6],
                    other.val[6],
                    self.val[7],
                    other.val[7],
                ],
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn slide<const SHIFT: usize>(self, next: Self) -> Self {
            let mut val = [0.0f32; SIMD_WIDTH];
            val[..SIMD_WIDTH - SHIFT].copy_from_slice(&self.val[SHIFT..]);
            val[SIMD_WIDTH - SHIFT..].copy_from_slice(&next.val[..SHIFT]);
            Self {
                val,
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        pub(crate) fn deinterleave(self, next: Self) -> (Self, Self) {
            let lanes = |parity: usize| Self {
                val: core::array::from_fn(|lane| {
                    let index = 2 * lane + parity;
                    *self.val.get(index).unwrap_or(&next.val[index % SIMD_WIDTH])
                }),
                _marker: PhantomData,
            };
            (lanes(0), lanes(1))
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        pub(crate) fn min(self, other: Self) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = super::min_f32(self.val[i], other.val[i]);
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }

        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        pub(crate) fn max(self, other: Self) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = super::max_f32(self.val[i], other.val[i]);
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }
    }

    impl<S: Simd> Add for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        fn add(self, rhs: Self) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = self.val[i] + rhs.val[i];
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }
    }

    impl<S: Simd> Sub for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        fn sub(self, rhs: Self) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = self.val[i] - rhs.val[i];
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }
    }

    impl<S: Simd> Mul for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        fn mul(self, rhs: Self) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = self.val[i] * rhs.val[i];
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }
    }

    impl<S: Simd> Add<f32> for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        fn add(self, rhs: f32) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = self.val[i] + rhs;
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }
    }

    impl<S: Simd> Mul<f32> for f32x8<S> {
        type Output = Self;
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[expect(clippy::needless_range_loop, reason = "fixed-width scalar SIMD lanes")]
        #[inline(always)]
        fn mul(self, rhs: f32) -> Self {
            let mut result = [0.0f32; SIMD_WIDTH];
            for i in 0..SIMD_WIDTH {
                result[i] = self.val[i] * rhs;
            }
            Self {
                val: result,
                _marker: PhantomData,
            }
        }
    }

    impl<S: Simd> AddAssign for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn add_assign(&mut self, rhs: Self) {
            for i in 0..SIMD_WIDTH {
                self.val[i] += rhs.val[i];
            }
        }
    }

    impl<S: Simd> SubAssign for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn sub_assign(&mut self, rhs: Self) {
            for i in 0..SIMD_WIDTH {
                self.val[i] -= rhs.val[i];
            }
        }
    }

    impl<S: Simd> MulAssign<f32> for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn mul_assign(&mut self, rhs: f32) {
            for i in 0..SIMD_WIDTH {
                self.val[i] *= rhs;
            }
        }
    }

    impl<S: Simd> DivAssign<f32> for f32x8<S> {
        #[expect(
            clippy::inline_always,
            reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
        )]
        #[inline(always)]
        fn div_assign(&mut self, rhs: f32) {
            for i in 0..SIMD_WIDTH {
                self.val[i] /= rhs;
            }
        }
    }

    /// Scalar fallback for SIMD dispatch.
    #[macro_export]
    macro_rules! simd_dispatch {
        ($level:expr, $simd:ident => $body:expr) => {{
            let _ = $level;
            let $simd = $crate::math::ScalarSimd;
            $body
        }};
    }

    pub(crate) use simd_dispatch as dispatch;
}

// Note that these polyfills can be very imprecise, but hopefully good enough
// for the vast majority of cases.

#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[inline(always)]
pub(crate) fn mul_add(a: f32, b: f32, c: f32) -> f32 {
    // The 9/7 lifting path must use one rounding step on every CPU target so
    // its output does not change with compile-time FMA availability. Hosted
    // builds use the hardware-backed intrinsic; no-std keeps the libm oracle.
    #[cfg(feature = "std")]
    let result = a.mul_add(b, c);
    #[cfg(not(feature = "std"))]
    let result = libm::fmaf(a, b, c);
    result
}

#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[cfg_attr(
    not(feature = "std"),
    expect(
        clippy::cast_possible_truncation,
        clippy::cast_precision_loss,
        reason = "the no-std floor polyfill converts through the truncated integer part"
    )
)]
#[inline(always)]
pub(crate) fn floor_f32(x: f32) -> f32 {
    #[cfg(feature = "std")]
    {
        x.floor()
    }
    #[cfg(not(feature = "std"))]
    {
        let xi = x as i32;
        let xf = xi as f32;
        if x < xf {
            xf - 1.0
        } else {
            xf
        }
    }
}

/// `floor((a + b) / 2)` with the sum rounded in `f32`, as 5/3 lifting needs.
#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[expect(
    clippy::manual_midpoint,
    reason = "5/3 lifting must round like its f32x8 lanes; f32::midpoint widens to f64"
)]
#[inline(always)]
pub(crate) fn floor_half_sum_f32(a: f32, b: f32) -> f32 {
    floor_f32((a + b) * 0.5)
}

#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[inline(always)]
pub(crate) fn round_f32(x: f32) -> f32 {
    #[cfg(feature = "std")]
    {
        x.round()
    }
    #[cfg(not(feature = "std"))]
    {
        if x >= 0.0 {
            floor_f32(x + 0.5)
        } else {
            -floor_f32(-x + 0.5)
        }
    }
}

#[expect(
    clippy::inline_always,
    reason = "the centered irreversible-output rounding boundary is shared by hot decode loops"
)]
#[inline(always)]
pub(crate) fn round_ties_even_then_add(sample: f32, addend: f32) -> f32 {
    #[cfg(feature = "std")]
    let rounded = sample.round_ties_even();
    #[cfg(not(feature = "std"))]
    let rounded = libm::roundevenf(sample);
    rounded + addend
}

#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[inline(always)]
pub(crate) fn log2_f32(x: f32) -> f32 {
    #[cfg(feature = "std")]
    {
        x.log2()
    }
    #[cfg(not(feature = "std"))]
    {
        libm::log2f(x)
    }
}

#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[inline(always)]
pub(crate) fn pow2i(exp: i32) -> f32 {
    // Build the power from its IEEE-754 bits so every exponent is exact, not
    // only those a 32-bit shift reaches: high-bit-depth quantization steps
    // need powers such as 2^38. Below the normal range the result is
    // subnormal, then zero; above it the result is infinity.
    if exp > 127 {
        f32::INFINITY
    } else if exp >= -126 {
        f32::from_bits((exp + 127).unsigned_abs() << 23)
    } else if exp >= -149 {
        f32::from_bits(1_u32 << (exp + 149).unsigned_abs())
    } else {
        0.0
    }
}

#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[inline(always)]
#[cfg(not(feature = "simd"))]
pub(crate) fn min_f32(a: f32, b: f32) -> f32 {
    #[cfg(feature = "std")]
    {
        a.min(b)
    }
    #[cfg(not(feature = "std"))]
    {
        if a < b {
            a
        } else {
            b
        }
    }
}

#[expect(
    clippy::inline_always,
    reason = "tiny SIMD and scalar wrappers must disappear inside transform loops"
)]
#[inline(always)]
#[cfg(not(feature = "simd"))]
pub(crate) fn max_f32(a: f32, b: f32) -> f32 {
    #[cfg(feature = "std")]
    {
        a.max(b)
    }
    #[cfg(not(feature = "std"))]
    {
        if a > b {
            a
        } else {
            b
        }
    }
}

#[cfg(not(feature = "simd"))]
pub(crate) use inner::{dispatch, f32x8, Level, ScalarSimd, Simd};
#[cfg(feature = "simd")]
pub(crate) use inner::{dispatch, f32x8, Level, Simd};

/// A wrapper around `Vec<f32>` that pads the vector to a multiple of `N` elements.
/// This allows SIMD operations to safely process the data without bounds checking
/// at the end of the buffer.
#[derive(Debug)]
pub(crate) struct SimdBuffer<const N: usize> {
    data: Vec<f32>,
    original_len: usize,
}

crate::move_only::assert_move_only!(SimdBuffer<8>);

#[derive(Debug, Clone, Copy)]
pub(crate) struct SimdBufferAllocationError;

#[cfg(test)]
std::thread_local! {
    static RETAINED_ZEROED_SAMPLES: core::cell::Cell<usize> = const { core::cell::Cell::new(0) };
}

#[cfg(test)]
pub(crate) fn reset_retained_zeroed_samples() {
    RETAINED_ZEROED_SAMPLES.set(0);
}

/// Samples of retained (reset, not freshly allocated) buffers zeroed on this
/// thread since the last reset.
#[cfg(test)]
pub(crate) fn retained_zeroed_samples() -> usize {
    RETAINED_ZEROED_SAMPLES.get()
}

#[cfg(test)]
fn note_retained_zeroed(samples: usize) {
    RETAINED_ZEROED_SAMPLES.set(RETAINED_ZEROED_SAMPLES.get() + samples);
}

impl<const N: usize> SimdBuffer<N> {
    /// Create an empty buffer without allocating.
    pub(crate) const fn empty() -> Self {
        Self {
            data: Vec::new(),
            original_len: 0,
        }
    }

    /// Create a new `SimdBuffer` from a `Vec<f32>`, padding it to a multiple of `N`.
    #[cfg(test)]
    pub(crate) fn new(mut data: Vec<f32>) -> Self {
        let original_len = data.len();
        let padded_len = Self::padded_len(original_len)
            .expect("an allocated vector length plus SIMD padding must fit usize");
        if padded_len > original_len {
            data.resize(padded_len, 0.0);
        }
        Self { data, original_len }
    }

    /// Create a new `SimdBuffer` filled with zeros using fallible allocation.
    pub(crate) fn try_zeros(
        original_len: usize,
    ) -> core::result::Result<Self, SimdBufferAllocationError> {
        let padded_len = Self::padded_len(original_len).ok_or(SimdBufferAllocationError)?;
        let mut data = Vec::new();
        data.try_reserve_exact(padded_len)
            .map_err(|_| SimdBufferAllocationError)?;
        data.resize(padded_len, 0.0);
        Ok(Self { data, original_len })
    }

    /// Create a zeroed buffer of `original_len` samples whose allocation also
    /// holds `capacity_len` samples, for callers that later fill it through
    /// [`Self::padded_scratch`].
    pub(crate) fn try_zeros_with_capacity(
        original_len: usize,
        capacity_len: usize,
    ) -> core::result::Result<Self, SimdBufferAllocationError> {
        let padded_len = Self::padded_len(original_len).ok_or(SimdBufferAllocationError)?;
        let capacity =
            Self::padded_len(capacity_len.max(original_len)).ok_or(SimdBufferAllocationError)?;
        let mut data = Vec::new();
        data.try_reserve_exact(capacity)
            .map_err(|_| SimdBufferAllocationError)?;
        data.resize(padded_len, 0.0);
        Ok(Self { data, original_len })
    }

    /// The first `len` samples of the allocation, which may extend past the
    /// logical length into retained capacity. Returns `None` when the
    /// allocation is smaller. Never reallocates. The logical samples keep their
    /// values; samples past them are zero the first time they are exposed.
    /// Callers must restore the layout with [`Self::compact_from_padding`].
    pub(crate) fn padded_scratch(&mut self, len: usize) -> Option<&mut [f32]> {
        if len > self.data.capacity() {
            return None;
        }
        if self.data.len() < len {
            self.data.resize(len, 0.0);
        }
        Some(&mut self.data[..len])
    }

    /// Move a `width` x `height` window at (`x`, `y`) of a row-major image
    /// `stride` samples wide, written through [`Self::padded_scratch`], to the
    /// front as the buffer's `width`-wide samples, passing each moved row to
    /// `finish_row`. Restores the logical length with zeroed SIMD padding.
    pub(crate) fn compact_from_padding(
        &mut self,
        stride: usize,
        (x, y): (usize, usize),
        (width, height): (usize, usize),
        mut finish_row: impl FnMut(&mut [f32]),
    ) -> core::result::Result<(), SimdBufferAllocationError> {
        let fits = x.checked_add(width).is_some_and(|end| end <= stride)
            && width.checked_mul(height) == Some(self.original_len)
            && y.checked_add(height)
                .and_then(|rows| rows.checked_mul(stride))
                .is_some_and(|end| end <= self.data.len());
        if !fits {
            return Err(SimdBufferAllocationError);
        }
        // Each destination row starts at or before its source row, so moving
        // rows front to back never overwrites a row still to be moved.
        for row in 0..height {
            let source = (y + row) * stride + x;
            let target = row * width;
            self.data.copy_within(source..source + width, target);
            finish_row(&mut self.data[target..target + width]);
        }
        let padded_len = Self::padded_len(self.original_len).ok_or(SimdBufferAllocationError)?;
        self.data.truncate(padded_len);
        self.data[self.original_len..].fill(0.0);
        Ok(())
    }

    /// Reset the logical length and zero every accessible sample while retaining
    /// an existing allocation whenever it is large enough.
    pub(crate) fn try_reset_zeros(
        &mut self,
        original_len: usize,
    ) -> core::result::Result<(), SimdBufferAllocationError> {
        let padded_len = Self::padded_len(original_len).ok_or(SimdBufferAllocationError)?;
        if padded_len > self.data.capacity() {
            return Err(SimdBufferAllocationError);
        }
        self.data.resize(padded_len, 0.0);
        self.data.fill(0.0);
        #[cfg(test)]
        note_retained_zeroed(padded_len);
        self.original_len = original_len;
        Ok(())
    }

    /// Reset the logical length for a decode that overwrites every sample in
    /// `[0, original_len)`, retaining an existing allocation whenever it is
    /// large enough. Those samples keep stale values until overwritten; only
    /// newly grown samples and the SIMD padding tail are zeroed.
    pub(crate) fn try_reset_for_overwrite(
        &mut self,
        original_len: usize,
    ) -> core::result::Result<(), SimdBufferAllocationError> {
        let padded_len = Self::padded_len(original_len).ok_or(SimdBufferAllocationError)?;
        if padded_len > self.data.capacity() {
            return Err(SimdBufferAllocationError);
        }
        // Growth and the tail together zero exactly [min(len, original_len), padded_len).
        #[cfg(test)]
        note_retained_zeroed(padded_len - self.data.len().min(original_len));
        self.data.resize(padded_len, 0.0);
        self.data[original_len..].fill(0.0);
        self.original_len = original_len;
        Ok(())
    }

    /// Exchange the logical samples with `samples`, which must hold exactly
    /// that many and room for the SIMD padding; `samples` receives the previous
    /// allocation. Returns `false`, changing nothing, when they do not fit.
    pub(crate) fn swap_samples(&mut self, samples: &mut Vec<f32>) -> bool {
        let Some(padded_len) = Self::padded_len(self.original_len) else {
            return false;
        };
        if samples.len() != self.original_len || samples.capacity() < padded_len {
            return false;
        }
        core::mem::swap(&mut self.data, samples);
        self.data.resize(padded_len, 0.0);
        true
    }

    /// Take `samples` as the logical samples of an unallocated buffer, when
    /// they have room for the SIMD padding; `samples` is left empty. Returns
    /// `false`, changing nothing, otherwise.
    pub(crate) fn adopt_samples(&mut self, samples: &mut Vec<f32>) -> bool {
        let Some(padded_len) = Self::padded_len(samples.len()) else {
            return false;
        };
        if self.data.capacity() != 0 || samples.capacity() < padded_len {
            return false;
        }
        self.original_len = samples.len();
        self.data = core::mem::take(samples);
        self.data.resize(padded_len, 0.0);
        true
    }

    /// Returns only the original (non-padded) data as an immutable slice.
    pub(crate) fn truncated(&self) -> &[f32] {
        &self.data[..self.original_len]
    }

    /// Allocated element capacity retained by this decode buffer.
    pub(crate) fn capacity(&self) -> usize {
        self.data.capacity()
    }

    /// Returns the length padded to a multiple of `N`
    pub(crate) fn padded_len(original_len: usize) -> Option<usize> {
        if N == 0 {
            return None;
        }
        let remainder = original_len % N;
        let padding = (N - remainder) % N;
        original_len.checked_add(padding)
    }
}

impl<const N: usize> core::ops::Deref for SimdBuffer<N> {
    type Target = [f32];

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

impl<const N: usize> core::ops::DerefMut for SimdBuffer<N> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.data
    }
}

#[cfg(test)]
mod simd_buffer_padding_tests {
    use super::{SimdBuffer, SIMD_WIDTH};

    #[test]
    fn a_window_written_into_padding_compacts_to_the_logical_samples() {
        // A 3 x 2 plane whose allocation also holds a 5 x 4 window.
        let mut buffer = SimdBuffer::<SIMD_WIDTH>::try_zeros_with_capacity(6, 20).expect("alloc");
        let capacity = buffer.capacity();
        let window = buffer.padded_scratch(20).expect("room for the window");
        for (sample, index) in window.iter_mut().zip(0_u8..) {
            *sample = f32::from(index);
        }
        buffer
            .compact_from_padding(5, (1, 1), (3, 2), |row| {
                for sample in row {
                    *sample += 0.5;
                }
            })
            .expect("compact");
        assert_eq!(buffer.truncated(), &[6.5, 7.5, 8.5, 11.5, 12.5, 13.5]);
        assert_eq!(
            buffer.len(),
            SimdBuffer::<SIMD_WIDTH>::padded_len(6).unwrap()
        );
        assert!(buffer[6..].iter().all(|&sample| sample == 0.0));
        assert_eq!(buffer.capacity(), capacity);
    }

    #[test]
    fn padding_is_never_grown_and_compaction_checks_the_window() {
        let mut buffer = SimdBuffer::<SIMD_WIDTH>::try_zeros(6).expect("alloc");
        assert!(buffer.padded_scratch(buffer.capacity() + 1).is_none());
        buffer.padded_scratch(6).expect("logical samples");
        assert!(buffer
            .compact_from_padding(3, (1, 0), (3, 2), |_| {})
            .is_err());
        assert!(buffer
            .compact_from_padding(3, (0, 0), (2, 2), |_| {})
            .is_err());
        assert!(buffer
            .compact_from_padding(3, (0, 0), (3, 2), |_| {})
            .is_ok());
    }
}

#[cfg(test)]
mod simd_operator_tests {
    use super::{dispatch, f32x8, Level, SIMD_WIDTH};

    #[test]
    fn assignment_operators_preserve_each_lane_for_active_simd_level() {
        let left = [1.0, -2.0, 3.5, -4.5, 5.25, -6.25, 7.75, -8.75];
        let right = [8.0, 7.0, -6.0, -5.0, 4.0, 3.0, -2.0, -1.0];
        let mut output = [0.0; SIMD_WIDTH];

        dispatch!(Level::new(), simd => {
            let mut value = f32x8::from_slice(simd, &left);
            value += f32x8::from_slice(simd, &right);
            value -= f32x8::splat(simd, 1.0);
            value *= 0.5;
            value /= 2.0;
            value.store(&mut output);
        });

        let expected: [f32; SIMD_WIDTH] =
            core::array::from_fn(|index| (left[index] + right[index] - 1.0) / 4.0);
        for (index, (actual, expected)) in output.into_iter().zip(expected).enumerate() {
            assert!(
                (actual - expected).abs() <= f32::EPSILON,
                "lane {index}: expected {expected}, got {actual}"
            );
        }
    }

    #[test]
    fn vector_mul_add_preserves_fused_rounding() {
        let a = [f32::from_bits(0x4526_ba09); SIMD_WIDTH];
        let b = [f32::from_bits(0xbf41_420c); SIMD_WIDTH];
        let c = [f32::from_bits(0x4470_e5c9); SIMD_WIDTH];
        let mut output = [0.0; SIMD_WIDTH];

        dispatch!(Level::new(), simd => {
            f32x8::from_slice(simd, &a)
                .mul_add(f32x8::from_slice(simd, &b), f32x8::from_slice(simd, &c))
                .store(&mut output);
        });

        assert!(output.iter().all(|value| value.to_bits() == 0xc483_47a5));
    }
}

#[cfg(test)]
mod floating_point_tests {
    use super::{mul_add, pow2i};

    #[test]
    fn pow2i_is_exact_across_the_whole_f32_exponent_range() {
        let mut expected = 1.0_f32;
        for exp in 0..=128 {
            assert_eq!(pow2i(exp).to_bits(), expected.to_bits(), "2^{exp}");
            expected *= 2.0;
        }
        let mut expected = 1.0_f32;
        for exp in (-150..=0).rev() {
            assert_eq!(pow2i(exp).to_bits(), expected.to_bits(), "2^{exp}");
            expected /= 2.0;
        }
        assert_eq!(pow2i(i32::MAX), f32::INFINITY);
        assert_eq!(pow2i(i32::MIN).to_bits(), 0.0_f32.to_bits());
    }
    #[cfg(feature = "std")]
    use j2k_codec_math::dwt;

    #[cfg(feature = "std")]
    fn assert_fused_mul_add_matches_libm(a: f32, b: f32, c: f32) {
        debug_assert!(a.is_finite() && b.is_finite() && c.is_finite());
        assert_eq!(
            mul_add(a, b, c).to_bits(),
            libm::fmaf(a, b, c).to_bits(),
            "fused multiply-add mismatch for a={a:e} ({:#010x}), b={b:e} ({:#010x}), c={c:e} ({:#010x})",
            a.to_bits(),
            b.to_bits(),
            c.to_bits()
        );
    }

    #[test]
    fn scalar_mul_add_has_one_rounding_step_on_every_target() {
        let value = mul_add(
            f32::from_bits(0x4526_ba09),
            f32::from_bits(0xbf41_420c),
            f32::from_bits(0x4470_e5c9),
        );

        assert_eq!(value.to_bits(), 0xc483_47a5);
    }

    #[cfg(feature = "std")]
    #[test]
    fn hosted_mul_add_matches_portable_fused_oracle_for_finite_inputs() {
        const BOUNDARY_VALUES: [f32; 26] = [
            0.0,
            -0.0,
            f32::from_bits(0x0000_0001),
            f32::from_bits(0x8000_0001),
            f32::from_bits(0x007f_ffff),
            f32::from_bits(0x807f_ffff),
            f32::MIN_POSITIVE,
            -f32::MIN_POSITIVE,
            0.5,
            -0.5,
            1.0,
            -1.0,
            2.0,
            -2.0,
            f32::MAX,
            f32::MIN,
            dwt::DWT97_ALPHA_F32,
            dwt::DWT97_BETA_F32,
            dwt::DWT97_GAMMA_F32,
            dwt::DWT97_DELTA_F32,
            dwt::DWT97_KAPPA_F32,
            dwt::DWT97_INV_KAPPA_F32,
            dwt::IDWT97_NEG_ALPHA_F32,
            dwt::IDWT97_NEG_BETA_F32,
            dwt::IDWT97_NEG_GAMMA_F32,
            dwt::IDWT97_NEG_DELTA_F32,
        ];

        for a in BOUNDARY_VALUES {
            for b in BOUNDARY_VALUES {
                for c in BOUNDARY_VALUES {
                    assert_fused_mul_add_matches_libm(a, b, c);
                }
            }
        }

        // A fixed generator makes this test reproducible while covering normal,
        // subnormal, cancellation, overflow, and signed-zero results broadly.
        let mut state = 0x243f_6a88_85a3_08d3_u64;
        let mut next_finite = || loop {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            let bits =
                u32::try_from(state & u64::from(u32::MAX)).expect("masked random state fits u32");
            let value = f32::from_bits(bits);
            if value.is_finite() {
                break value;
            }
        };

        for _ in 0..100_000 {
            let a = next_finite();
            let b = next_finite();
            let c = next_finite();
            assert_fused_mul_add_matches_libm(a, b, c);
        }
    }
}

#[cfg(test)]
mod integer_tests {
    use super::{bit_width_u32, bit_width_u64, ceil_log2_u32, SimdBuffer};

    #[test]
    fn integer_bit_ranges_cover_type_boundaries() {
        assert_eq!(bit_width_u32(0), 0);
        assert_eq!(bit_width_u32(1), 1);
        assert_eq!(bit_width_u32(u32::MAX), 32);
        assert_eq!(bit_width_u64(0), 0);
        assert_eq!(bit_width_u64(1_u64 << 63), 64);
        assert_eq!(bit_width_u64(u64::MAX), 64);
    }

    #[test]
    fn simd_buffer_fallible_zero_allocation_uses_only_required_padding() {
        let aligned = SimdBuffer::<8>::try_zeros(16).unwrap();
        assert_eq!(aligned.len(), 16);
        assert_eq!(aligned.truncated(), &[0.0; 16]);

        let unaligned = SimdBuffer::<8>::try_zeros(17).unwrap();
        assert_eq!(unaligned.len(), 24);
        assert_eq!(unaligned.truncated(), &[0.0; 17]);

        assert!(SimdBuffer::<8>::try_zeros(usize::MAX).is_err());
        assert!(SimdBuffer::<0>::try_zeros(1).is_err());
    }

    #[test]
    fn ceil_log2_covers_zero_and_full_u32_domain() {
        assert_eq!(ceil_log2_u32(0), 0);
        assert_eq!(ceil_log2_u32(1), 0);
        assert_eq!(ceil_log2_u32(2), 1);
        assert_eq!(ceil_log2_u32(3), 2);
        assert_eq!(ceil_log2_u32(u32::MAX), 32);
    }
}