maroontree 0.1.3

/*
 * Copyright (c) Radzivon Bartoshyk 7/2026. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 * 1.  Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 *
 * 2.  Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3.  Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#![allow(clippy::needless_range_loop)]

use crate::dct::{
    SQRT2, WC4_0, WC4_1, WC8_0, WC8_1, WC8_2, WC8_3, WC16_0, WC16_1, WC16_2, WC16_3, WC16_4,
    WC16_5, WC16_6, WC16_7, WC32,
};
#[cfg(target_arch = "x86")]
use std::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;
use std::mem::MaybeUninit;

#[derive(Clone, Copy)]
struct I32x8(__m256i);

#[inline]
#[target_feature(enable = "avx2")]
fn zero() -> I32x8 {
    I32x8(_mm256_setzero_si256())
}

#[inline]
#[target_feature(enable = "avx2")]
fn splat(v: i32) -> __m256i {
    _mm256_set1_epi32(v)
}

#[inline]
#[target_feature(enable = "avx2")]
fn load_i32x8(ptr: *const i32) -> I32x8 {
    unsafe { I32x8(_mm256_loadu_si256(ptr.cast::<__m256i>())) }
}

#[inline]
#[target_feature(enable = "avx2")]
fn store_i32x8(ptr: *mut i32, v: I32x8) {
    unsafe { _mm256_storeu_si256(ptr.cast::<__m256i>(), v.0) };
}

#[inline]
#[target_feature(enable = "avx2")]
fn sar_epi64_16(v: __m256i) -> __m256i {
    let sign = _mm256_cmpgt_epi64(_mm256_setzero_si256(), v);
    _mm256_or_si256(_mm256_srli_epi64::<16>(v), _mm256_slli_epi64::<48>(sign))
}

#[inline]
#[target_feature(enable = "avx2")]
fn mul_q16_epi32(v: __m256i, c: __m256i) -> __m256i {
    let even = sar_epi64_16(_mm256_mul_epi32(v, c));
    let odd_v = _mm256_srli_epi64::<32>(v);
    let odd_c = _mm256_srli_epi64::<32>(c);
    let odd = _mm256_slli_epi64::<32>(sar_epi64_16(_mm256_mul_epi32(odd_v, odd_c)));
    _mm256_blend_epi32::<0b1010_1010>(even, odd)
}

#[inline]
fn quant_flat<const N: usize>(coeffs: &[i32; N], dc_q: i32, ac_q: i32, out: &mut [i32; N]) {
    let mq = |a: i32, b: i32| {
        let prod = (a as i64) * (b as i64);
        let mag = prod.unsigned_abs();
        if mag < 65536 {
            return 0;
        }
        let lvl = ((mag + 32768) >> 16) as i32;
        if prod >= 0 { lvl } else { -lvl }
    };
    out[0] = mq(coeffs[0], dc_q);
    for i in 1..N {
        out[i] = mq(coeffs[i], ac_q);
    }
}

impl I32x8 {
    #[inline]
    #[target_feature(enable = "avx2")]
    fn add(self, rhs: I32x8) -> I32x8 {
        I32x8(_mm256_add_epi32(self.0, rhs.0))
    }

    #[inline]
    #[target_feature(enable = "avx2")]
    fn sub(self, rhs: I32x8) -> I32x8 {
        I32x8(_mm256_sub_epi32(self.0, rhs.0))
    }

    #[inline]
    #[target_feature(enable = "avx2")]
    fn muls_q16(self, coeff: i32) -> I32x8 {
        I32x8(mul_q16_epi32(self.0, splat(coeff)))
    }

    #[inline]
    #[target_feature(enable = "avx2")]
    fn fma_sqrt2(self, rhs: I32x8) -> I32x8 {
        self.muls_q16(SQRT2).add(rhs)
    }

    #[inline]
    #[target_feature(enable = "avx2")]
    fn shr<const N: i32>(self) -> I32x8 {
        I32x8(_mm256_srai_epi32::<N>(self.0))
    }

    #[inline]
    #[target_feature(enable = "avx2")]
    fn shl<const N: i32>(self) -> I32x8 {
        I32x8(_mm256_slli_epi32::<N>(self.0))
    }

    #[inline]
    #[target_feature(enable = "avx2")]
    fn shr_round<const N: i32>(self) -> I32x8 {
        I32x8(_mm256_srai_epi32::<N>(_mm256_add_epi32(
            self.0,
            splat(1 << (N - 1)),
        )))
    }
}

#[inline]
#[target_feature(enable = "avx2")]
fn transpose_8x8_i32(c: &mut [I32x8; 8]) {
    let r0 = c[0].0;
    let r1 = c[1].0;
    let r2 = c[2].0;
    let r3 = c[3].0;
    let r4 = c[4].0;
    let r5 = c[5].0;
    let r6 = c[6].0;
    let r7 = c[7].0;

    let t0 = _mm256_unpacklo_epi32(r0, r1);
    let t1 = _mm256_unpackhi_epi32(r0, r1);
    let t2 = _mm256_unpacklo_epi32(r2, r3);
    let t3 = _mm256_unpackhi_epi32(r2, r3);
    let t4 = _mm256_unpacklo_epi32(r4, r5);
    let t5 = _mm256_unpackhi_epi32(r4, r5);
    let t6 = _mm256_unpacklo_epi32(r6, r7);
    let t7 = _mm256_unpackhi_epi32(r6, r7);

    let u0 = _mm256_unpacklo_epi64(t0, t2);
    let u1 = _mm256_unpackhi_epi64(t0, t2);
    let u2 = _mm256_unpacklo_epi64(t1, t3);
    let u3 = _mm256_unpackhi_epi64(t1, t3);
    let u4 = _mm256_unpacklo_epi64(t4, t6);
    let u5 = _mm256_unpackhi_epi64(t4, t6);
    let u6 = _mm256_unpacklo_epi64(t5, t7);
    let u7 = _mm256_unpackhi_epi64(t5, t7);

    c[0] = I32x8(_mm256_permute2x128_si256::<0x20>(u0, u4));
    c[1] = I32x8(_mm256_permute2x128_si256::<0x20>(u1, u5));
    c[2] = I32x8(_mm256_permute2x128_si256::<0x20>(u2, u6));
    c[3] = I32x8(_mm256_permute2x128_si256::<0x20>(u3, u7));
    c[4] = I32x8(_mm256_permute2x128_si256::<0x31>(u0, u4));
    c[5] = I32x8(_mm256_permute2x128_si256::<0x31>(u1, u5));
    c[6] = I32x8(_mm256_permute2x128_si256::<0x31>(u2, u6));
    c[7] = I32x8(_mm256_permute2x128_si256::<0x31>(u3, u7));
}

#[inline]
#[target_feature(enable = "avx2")]
fn dct1d_4_v_i32(c: &mut [I32x8; 4]) {
    let t0 = c[0].add(c[3]);
    let t1 = c[1].add(c[2]);
    let sum = t0.add(t1);
    let diff = t0.sub(t1);

    let t2 = c[0].sub(c[3]).muls_q16(WC4_0);
    let t3 = c[1].sub(c[2]).muls_q16(WC4_1);
    let t2p = t2.add(t3);
    let t3p = t2.sub(t3);
    let t2pp = t2p.fma_sqrt2(t3p);

    c[0] = sum;
    c[1] = t2pp;
    c[2] = diff;
    c[3] = t3p;
}

#[inline]
#[target_feature(enable = "avx2")]
fn dct1d_8_v_i32(c: &mut [I32x8; 8]) {
    let mut evens = [
        c[0].add(c[7]),
        c[1].add(c[6]),
        c[2].add(c[5]),
        c[3].add(c[4]),
    ];
    dct1d_4_v_i32(&mut evens);

    let mut odds = [
        c[0].sub(c[7]).muls_q16(WC8_0),
        c[1].sub(c[6]).muls_q16(WC8_1),
        c[2].sub(c[5]).muls_q16(WC8_2),
        c[3].sub(c[4]).muls_q16(WC8_3),
    ];
    dct1d_4_v_i32(&mut odds);

    odds[0] = odds[0].fma_sqrt2(odds[1]);
    odds[1] = odds[1].add(odds[2]);
    odds[2] = odds[2].add(odds[3]);

    c[0] = evens[0];
    c[1] = odds[0];
    c[2] = evens[1];
    c[3] = odds[1];
    c[4] = evens[2];
    c[5] = odds[2];
    c[6] = evens[3];
    c[7] = odds[3];
}

#[inline]
#[target_feature(enable = "avx2")]
fn dct1d_16_v_i32(c: &mut [I32x8; 16]) {
    let mut evens = [
        c[0].add(c[15]),
        c[1].add(c[14]),
        c[2].add(c[13]),
        c[3].add(c[12]),
        c[4].add(c[11]),
        c[5].add(c[10]),
        c[6].add(c[9]),
        c[7].add(c[8]),
    ];
    let mut odds = [
        c[0].sub(c[15]).muls_q16(WC16_0),
        c[1].sub(c[14]).muls_q16(WC16_1),
        c[2].sub(c[13]).muls_q16(WC16_2),
        c[3].sub(c[12]).muls_q16(WC16_3),
        c[4].sub(c[11]).muls_q16(WC16_4),
        c[5].sub(c[10]).muls_q16(WC16_5),
        c[6].sub(c[9]).muls_q16(WC16_6),
        c[7].sub(c[8]).muls_q16(WC16_7),
    ];

    dct1d_8_v_i32(&mut evens);
    dct1d_8_v_i32(&mut odds);

    odds[0] = odds[0].fma_sqrt2(odds[1]);
    odds[1] = odds[1].add(odds[2]);
    odds[2] = odds[2].add(odds[3]);
    odds[3] = odds[3].add(odds[4]);
    odds[4] = odds[4].add(odds[5]);
    odds[5] = odds[5].add(odds[6]);
    odds[6] = odds[6].add(odds[7]);

    c[0] = evens[0];
    c[1] = odds[0];
    c[2] = evens[1];
    c[3] = odds[1];
    c[4] = evens[2];
    c[5] = odds[2];
    c[6] = evens[3];
    c[7] = odds[3];
    c[8] = evens[4];
    c[9] = odds[4];
    c[10] = evens[5];
    c[11] = odds[5];
    c[12] = evens[6];
    c[13] = odds[6];
    c[14] = evens[7];
    c[15] = odds[7];
}

#[inline]
#[target_feature(enable = "avx2")]
fn dct1d_32_v_i32(c: &mut [I32x8; 32]) {
    let mut evens = std::array::from_fn::<I32x8, 16, _>(|i| c[i].add(c[31 - i]));
    let mut odds = std::array::from_fn::<I32x8, 16, _>(|i| c[i].sub(c[31 - i]));

    dct1d_16_v_i32(&mut evens);
    for i in 0..16 {
        odds[i] = odds[i].muls_q16(WC32[i]);
    }
    dct1d_16_v_i32(&mut odds);

    odds[0] = odds[0].fma_sqrt2(odds[1]);
    odds[1] = odds[1].add(odds[2]);
    odds[2] = odds[2].add(odds[3]);
    odds[3] = odds[3].add(odds[4]);
    odds[4] = odds[4].add(odds[5]);
    odds[5] = odds[5].add(odds[6]);
    odds[6] = odds[6].add(odds[7]);
    odds[7] = odds[7].add(odds[8]);
    odds[8] = odds[8].add(odds[9]);
    odds[9] = odds[9].add(odds[10]);
    odds[10] = odds[10].add(odds[11]);
    odds[11] = odds[11].add(odds[12]);
    odds[12] = odds[12].add(odds[13]);
    odds[13] = odds[13].add(odds[14]);
    odds[14] = odds[14].add(odds[15]);

    for i in 0..16 {
        c[2 * i] = evens[i];
        c[2 * i + 1] = odds[i];
    }
}

#[inline]
#[target_feature(enable = "avx2")]
fn load8_i32(ptr: &[i32], stride: usize) -> [I32x8; 8] {
    std::array::from_fn(|y| load_i32x8(unsafe { ptr.as_ptr().add(y * stride) }))
}

#[inline]
#[target_feature(enable = "avx2")]
fn load16_i32(ptr: &[i32], stride: usize) -> [I32x8; 16] {
    std::array::from_fn(|y| load_i32x8(unsafe { ptr.as_ptr().add(y * stride) }))
}

#[inline]
#[target_feature(enable = "avx2")]
fn load32_i32(ptr: &[i32], stride: usize) -> [I32x8; 32] {
    std::array::from_fn(|y| load_i32x8(unsafe { ptr.as_ptr().add(y * stride) }))
}

#[target_feature(enable = "avx2")]
pub(crate) fn dct8x8_avx2_coeffs(input: &[i32; 64]) -> [i32; 64] {
    let mut cols = load8_i32(input, 8);
    dct1d_8_v_i32(&mut cols);
    transpose_8x8_i32(&mut cols);
    dct1d_8_v_i32(&mut cols);

    let mut out = MaybeUninit::<[i32; 64]>::uninit();
    for (k, col) in cols.iter().copied().enumerate() {
        unsafe { store_i32x8((out.as_mut_ptr() as *mut i32).add(k * 8), col) };
    }
    unsafe { out.assume_init() }
}

#[target_feature(enable = "avx2")]
pub(crate) fn dct16x16_avx2_coeffs(input: &[i32; 256]) -> [i32; 256] {
    let mut c_l = load16_i32(input, 16);
    let mut c_r = load16_i32(&input[8..], 16);

    dct1d_16_v_i32(&mut c_l);
    dct1d_16_v_i32(&mut c_r);

    let mut top_l: [I32x8; 8] = c_l[..8].try_into().unwrap();
    let mut bot_l: [I32x8; 8] = c_l[8..16].try_into().unwrap();
    let mut top_r: [I32x8; 8] = c_r[..8].try_into().unwrap();
    let mut bot_r: [I32x8; 8] = c_r[8..16].try_into().unwrap();

    transpose_8x8_i32(&mut top_l);
    transpose_8x8_i32(&mut bot_l);
    transpose_8x8_i32(&mut top_r);
    transpose_8x8_i32(&mut bot_r);

    let mut d_a = [zero(); 16];
    let mut d_b = [zero(); 16];
    d_a[..8].copy_from_slice(&top_l);
    d_a[8..16].copy_from_slice(&top_r);
    d_b[..8].copy_from_slice(&bot_l);
    d_b[8..16].copy_from_slice(&bot_r);

    dct1d_16_v_i32(&mut d_a);
    dct1d_16_v_i32(&mut d_b);

    let mut out = MaybeUninit::<[i32; 256]>::uninit();
    for u in 0..16usize {
        unsafe {
            let dst = (out.as_mut_ptr() as *mut i32).add(u * 16);
            store_i32x8(dst, d_a[u].shr::<1>());
            store_i32x8(dst.add(8), d_b[u].shr::<1>());
        }
    }
    unsafe { out.assume_init() }
}

#[target_feature(enable = "avx2")]
pub(crate) fn dct16x16_avx2_i32(input: &mut [i32; 256], dc_q: i32, ac_q: i32) {
    let coeffs = dct16x16_avx2_coeffs(input);
    quant_flat(&coeffs, dc_q, ac_q, input);
}

#[target_feature(enable = "avx2")]
pub(crate) fn dct32x32_avx2_coeffs(input: &[i32; 1024]) -> [i32; 1024] {
    let mut tmp_u = MaybeUninit::<[i32; 1024]>::uninit();

    for group in 0..4usize {
        let col_start = group * 8;
        let mut cols = load32_i32(&input[col_start..], 32);
        for c in cols.iter_mut() {
            *c = c.shl::<6>();
        }
        dct1d_32_v_i32(&mut cols);
        for v in 0..32usize {
            unsafe {
                store_i32x8(
                    (tmp_u.as_mut_ptr() as *mut i32).add(v * 32 + col_start),
                    cols[v],
                );
            }
        }
    }

    let tmp = unsafe { tmp_u.assume_init() };
    let mut out = MaybeUninit::<[i32; 1024]>::uninit();

    for group in 0..4usize {
        let row_start = group * 8;
        let base_off = row_start * 32;

        let mut seg_a = load8_i32(&tmp[base_off..], 32);
        let mut seg_b = load8_i32(&tmp[base_off + 8..], 32);
        let mut seg_c = load8_i32(&tmp[base_off + 16..], 32);
        let mut seg_d = load8_i32(&tmp[base_off + 24..], 32);

        transpose_8x8_i32(&mut seg_a);
        transpose_8x8_i32(&mut seg_b);
        transpose_8x8_i32(&mut seg_c);
        transpose_8x8_i32(&mut seg_d);

        let mut rows = [zero(); 32];
        rows[..8].copy_from_slice(&seg_a);
        rows[8..16].copy_from_slice(&seg_b);
        rows[16..24].copy_from_slice(&seg_c);
        rows[24..32].copy_from_slice(&seg_d);

        dct1d_32_v_i32(&mut rows);
        for u in 0..32usize {
            unsafe {
                store_i32x8(
                    (out.as_mut_ptr() as *mut i32).add(u * 32 + row_start),
                    rows[u].shr_round::<8>(),
                );
            }
        }
    }

    unsafe { out.assume_init() }
}

#[target_feature(enable = "avx2")]
pub(crate) fn dct32x32_avx2_i32(input: &mut [i32; 1024], dc_q: i32, ac_q: i32) {
    let coeffs = dct32x32_avx2_coeffs(input);
    quant_flat(&coeffs, dc_q, ac_q, input);
}

#[target_feature(enable = "avx2")]
pub(crate) fn dct8x16_avx2_coeffs(input: &[i32; 128]) -> [i32; 128] {
    let mut rows = load16_i32(input, 8);
    dct1d_16_v_i32(&mut rows);

    let mut a: [I32x8; 8] = rows[..8].try_into().unwrap();
    let mut b: [I32x8; 8] = rows[8..16].try_into().unwrap();
    transpose_8x8_i32(&mut a);
    transpose_8x8_i32(&mut b);
    dct1d_8_v_i32(&mut a);
    dct1d_8_v_i32(&mut b);

    let mut out = MaybeUninit::<[i32; 128]>::uninit();
    for fx in 0..8usize {
        unsafe {
            let dst = (out.as_mut_ptr() as *mut i32).add(fx * 16);
            store_i32x8(dst, a[fx].muls_q16(46341));
            store_i32x8(dst.add(8), b[fx].muls_q16(46341));
        }
    }
    unsafe { out.assume_init() }
}

#[target_feature(enable = "avx2")]
pub(crate) fn dct8x16_avx2_i32(input: &mut [i32; 128], dc_q: i32, ac_q: i32) {
    let coeffs = dct8x16_avx2_coeffs(input);
    quant_flat(&coeffs, dc_q, ac_q, input);
}

#[cfg(test)]
mod avx2_vs_scalar {
    use super::*;
    use crate::dct::{
        dct8x8_coeffs, dct8x16_coeffs, dct8x16_i32_scalar, dct16x16_coeffs, dct16x16_scalar,
        dct32x32_coeffs, dct32x32_scalar,
    };

    const QUANT_PAIRS: &[(i32, i32)] = &[(65536, 65536), (65536, 46341), (32768, 32768)];

    fn lcg(state: &mut u32) -> i32 {
        *state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
        ((*state >> 16) as i32 & 0x3ff) - 512
    }

    fn fill_lcg(buf: &mut [i32], seed: u32, mask: i32) {
        let mut s = seed;
        for v in buf.iter_mut() {
            *v = (lcg(&mut s) & mask) - ((mask + 1) >> 1);
        }
    }

    #[test]
    fn dct8x8_coeffs_match_scalar() {
        if !std::is_x86_feature_detected!("avx2") {
            return;
        }
        for seed in [0x1234_5678, 0xdead_beef, 0xabcd_1234] {
            let mut input = [0i32; 64];
            fill_lcg(&mut input, seed, 0xff);
            let s = dct8x8_coeffs(&input);
            let a = unsafe { dct8x8_avx2_coeffs(&input) };
            assert_eq!(s, a, "8x8 coeff mismatch seed={seed:#x}");
        }
    }

    #[test]
    fn dct16x16_coeffs_and_levels_match_scalar() {
        if !std::is_x86_feature_detected!("avx2") {
            return;
        }
        for seed in [0x1234_5678, 0xdead_beef] {
            let mut input = [0i32; 256];
            fill_lcg(&mut input, seed, 0x1ff);
            assert_eq!(
                dct16x16_coeffs(&input),
                unsafe { dct16x16_avx2_coeffs(&input) },
                "16x16 coeff mismatch seed={seed:#x}"
            );
            for &(dc_q, ac_q) in QUANT_PAIRS {
                let mut s = input;
                let mut a = input;
                dct16x16_scalar(&mut s, dc_q, ac_q);
                unsafe { dct16x16_avx2_i32(&mut a, dc_q, ac_q) };
                assert_eq!(
                    s, a,
                    "16x16 level mismatch seed={seed:#x} dc_q={dc_q} ac_q={ac_q}"
                );
            }
        }
    }

    #[test]
    fn dct32x32_coeffs_and_levels_match_scalar() {
        if !std::is_x86_feature_detected!("avx2") {
            return;
        }
        for seed in [0x1234_5678, 0xdead_beef] {
            let mut input = [0i32; 1024];
            fill_lcg(&mut input, seed, 0x7f);
            assert_eq!(
                dct32x32_coeffs(&input),
                unsafe { dct32x32_avx2_coeffs(&input) },
                "32x32 coeff mismatch seed={seed:#x}"
            );
            for &(dc_q, ac_q) in QUANT_PAIRS {
                let mut s = input;
                let mut a = input;
                dct32x32_scalar(&mut s, dc_q, ac_q);
                unsafe { dct32x32_avx2_i32(&mut a, dc_q, ac_q) };
                assert_eq!(
                    s, a,
                    "32x32 level mismatch seed={seed:#x} dc_q={dc_q} ac_q={ac_q}"
                );
            }
        }
    }

    #[test]
    fn dct8x16_coeffs_and_levels_match_scalar() {
        if !std::is_x86_feature_detected!("avx2") {
            return;
        }
        for seed in [0x1234_5678, 0xdead_beef] {
            let mut input = [0i32; 128];
            fill_lcg(&mut input, seed, 0x1ff);
            assert_eq!(
                dct8x16_coeffs(&input),
                unsafe { dct8x16_avx2_coeffs(&input) },
                "8x16 coeff mismatch seed={seed:#x}"
            );
            for &(dc_q, ac_q) in QUANT_PAIRS {
                let mut s = input;
                let mut a = input;
                dct8x16_i32_scalar(&mut s, dc_q, ac_q);
                unsafe { dct8x16_avx2_i32(&mut a, dc_q, ac_q) };
                assert_eq!(
                    s, a,
                    "8x16 level mismatch seed={seed:#x} dc_q={dc_q} ac_q={ac_q}"
                );
            }
        }
    }
}