pxdct 0.3.2

/*
 * // Copyright (c) Radzivon Bartoshyk 1/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.
 */
use crate::avx::dct4::mixed_radix3f::dct4_radix_n_rotation_twiddles_avxf;
use crate::avx::storef::AvxStoreF;
use crate::avx::util::{define_avx_butterfly, fma};
use crate::bidirectional::{BidirectionalStore, InPlaceStore};
use crate::util::{DctConstants, DctSample};
use crate::{PxdctError, PxdctExecutor};
use num_traits::AsPrimitive;

#[derive(Debug, Clone)]
pub(crate) struct AvxDct4Butterfly3<T> {
    _phantom_data: std::marker::PhantomData<T>,
}

impl<T: DctSample> Default for AvxDct4Butterfly3<T>
where
    f64: AsPrimitive<T>,
{
    fn default() -> Self {
        Self {
            _phantom_data: std::marker::PhantomData {},
        }
    }
}

impl<T: DctSample> AvxDct4Butterfly3<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    pub(crate) fn exec<S: BidirectionalStore<T>>(&self, data: &mut S) {
        // DCT-IV Radix-Q for Q = 3
        let identity = T::FRAC_1_SQRT_2;
        let a = data[1] * identity;
        let c = (data[0] + data[2]) * identity;
        let s = (data[0] - data[2]) * identity;

        let u0 = fma(c, T::SQRT_3_OVER_2, s * T::HALF);
        let mut v0 = fma(c, T::HALF, -s * T::SQRT_3_OVER_2);

        let mut u1 = u0 * T::HALF;
        v0 *= T::SQRT_3_OVER_2;
        u1 = u1 - a;

        data[0] = u0 + a;
        data[1] = u1 - v0;
        data[2] = v0 + u1;
    }
}

define_avx_butterfly!(AvxDct4Butterfly3, 3);

#[derive(Debug, Clone)]
pub(crate) struct AvxDct4Butterfly9<T> {
    _phantom: std::marker::PhantomData<T>,
}

impl<T: DctSample> Default for AvxDct4Butterfly9<T>
where
    f64: AsPrimitive<T>,
{
    fn default() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T: DctSample> AvxDct4Butterfly9<T> {
    #[inline(always)]
    pub(crate) fn exec<S: BidirectionalStore<T>>(&self, data: &mut S) {
        // This is autogenerated Radix-Q where Q = 9.
        let a_buffer = [data[4] * T::FRAC_1_SQRT_2];
        let c_buffer = [
            (data[0] + data[8]) * T::FRAC_1_SQRT_2,
            (data[1] + data[7]) * T::FRAC_1_SQRT_2,
            (data[2] + data[6]) * T::FRAC_1_SQRT_2,
            (data[3] + data[5]) * T::FRAC_1_SQRT_2,
        ];
        let s_buffer = [
            (data[0] - data[8]) * T::FRAC_1_SQRT_2,
            (data[1] - data[7]) * T::FRAC_1_SQRT_2,
            (data[2] - data[6]) * T::FRAC_1_SQRT_2,
            (data[3] - data[5]) * T::FRAC_1_SQRT_2,
        ];
        let c0 = c_buffer[0];
        let s0 = s_buffer[0];
        let s0_tw0 = s0 * T::D4_R9_ROT_TWIDDLE_1;
        let s0_tw1 = -s0 * T::D4_R9_ROT_TWIDDLE_0;
        let iq0 = fma(c0, T::D4_R9_ROT_TWIDDLE_0, s0_tw0);
        let siq0 = fma(c0, T::D4_R9_ROT_TWIDDLE_1, s0_tw1);
        let mut u0 = iq0;
        let mut u1 = u0;
        let mut v0 = siq0;
        u1 *= T::D4_R9_ROT_TWIDDLE_2;
        v0 *= T::D4_R9_ROT_TWIDDLE_3;
        let c1 = c_buffer[1];
        let s1 = s_buffer[1];
        let iq1 = fma(c1, T::D4_R9_ROT_TWIDDLE_4, s1 * T::HALF);
        let siq1 = fma(c1, T::HALF, -s1 * T::D4_R9_ROT_TWIDDLE_4);
        u0 += iq1;
        u1 = fma(iq1, T::HALF, u1);
        v0 = fma(siq1, T::D4_R9_ROT_TWIDDLE_4, v0);
        let c2 = c_buffer[2];
        let s2 = s_buffer[2];
        let iq2 = fma(c2, T::D4_R9_ROT_TWIDDLE_2, s2 * T::D4_R9_ROT_TWIDDLE_3);
        let siq2 = fma(c2, T::D4_R9_ROT_TWIDDLE_3, -s2 * T::D4_R9_ROT_TWIDDLE_2);
        u0 += iq2;
        u1 = fma(iq2, T::D4_R9_ROT_TWIDDLE_5, u1);
        v0 = fma(siq2, T::D4_R9_ROT_TWIDDLE_6, v0);
        let c3 = c_buffer[3];
        let s3 = s_buffer[3];
        let iq3 = fma(c3, T::D4_R9_ROT_TWIDDLE_6, s3 * -T::D4_R9_ROT_TWIDDLE_5);
        let siq3 = fma(c3, -T::D4_R9_ROT_TWIDDLE_5, -s3 * T::D4_R9_ROT_TWIDDLE_6);
        u0 += iq3;
        u1 = fma(iq3, -T::D4_R9_ROT_TWIDDLE_0, u1);
        v0 = fma(siq3, T::D4_R9_ROT_TWIDDLE_1, v0);
        let a0 = a_buffer[0];
        u0 += a0;
        u1 = u1 - a0;
        data[0] = u0;
        data[1] = u1 - v0;
        data[2] = u1 + v0;
        let mut u2 = iq0;
        let mut v2 = siq0;
        u2 *= T::D4_R9_ROT_TWIDDLE_0;
        v2 *= T::D4_R9_ROT_TWIDDLE_1;
        u2 = fma(iq1, -T::HALF, u2);
        v2 = fma(siq1, T::D4_R9_ROT_TWIDDLE_4, v2);
        u2 = fma(iq2, -T::D4_R9_ROT_TWIDDLE_2, u2);
        v2 = fma(siq2, -T::D4_R9_ROT_TWIDDLE_3, v2);
        u2 = fma(iq3, -T::D4_R9_ROT_TWIDDLE_5, u2);
        v2 = fma(siq3, -T::D4_R9_ROT_TWIDDLE_6, v2);
        u2 += a0;
        data[3] = u2 - v2;
        data[4] = u2 + v2;
        let mut u3 = iq0;
        let mut v3 = siq0;
        u3 *= T::HALF;
        v3 *= T::D4_R9_ROT_TWIDDLE_4;
        u3 = u3 - iq1;
        u3 = fma(iq2, T::HALF, u3);
        v3 = fma(siq2, -T::D4_R9_ROT_TWIDDLE_4, v3);
        u3 = fma(iq3, T::HALF, u3);
        v3 = fma(siq3, T::D4_R9_ROT_TWIDDLE_4, v3);
        u3 = u3 - a0;
        data[5] = u3 - v3;
        data[6] = u3 + v3;
        let mut u4 = iq0;
        let mut v4 = siq0;
        u4 *= -T::D4_R9_ROT_TWIDDLE_5;
        v4 *= T::D4_R9_ROT_TWIDDLE_6;
        u4 = fma(iq1, -T::HALF, u4);
        v4 = fma(siq1, -T::D4_R9_ROT_TWIDDLE_4, v4);
        u4 = fma(iq2, T::D4_R9_ROT_TWIDDLE_0, u4);
        v4 = fma(siq2, T::D4_R9_ROT_TWIDDLE_1, v4);
        u4 = fma(iq3, -T::D4_R9_ROT_TWIDDLE_2, u4);
        v4 = fma(siq3, -T::D4_R9_ROT_TWIDDLE_3, v4);
        u4 += a0;
        data[7] = u4 - v4;
        data[8] = u4 + v4;
    }
}

define_avx_butterfly!(AvxDct4Butterfly9, 9);

#[derive(Debug, Clone)]
pub(crate) struct AvxDct4Butterfly27f {
    twiddle: [AvxStoreF; 4],
    bf9: AvxDct4Butterfly9<f32>,
}

impl Default for AvxDct4Butterfly27f {
    fn default() -> Self {
        Self {
            twiddle: unsafe {
                dct4_radix_n_rotation_twiddles_avxf(3, 9, 27)
                    .try_into()
                    .unwrap()
            },
            bf9: AvxDct4Butterfly9::default(),
        }
    }
}

impl AvxDct4Butterfly27f {
    #[inline]
    #[target_feature(enable = "avx2", enable = "fma")]
    pub(crate) fn exec<S: BidirectionalStore<f32>>(&self, data: &mut S) {
        let mut a_buffer = [
            data[1], data[4], data[7], data[10], data[13], data[16], data[19], data[22], data[25],
        ];
        let mut c_buffer = [
            data[0] + data[2],
            data[3] + data[5],
            data[6] + data[8],
            data[9] + data[11],
            data[12] + data[14],
            data[15] + data[17],
            data[18] + data[20],
            data[21] + data[23],
            data[24] + data[26],
        ];
        let mut s_buffer = [
            data[0] - data[2],
            data[5] - data[3],
            data[6] - data[8],
            data[11] - data[9],
            data[12] - data[14],
            data[17] - data[15],
            data[18] - data[20],
            data[23] - data[21],
            data[24] - data[26],
        ];

        self.bf9.exec(&mut InPlaceStore::new(&mut a_buffer));
        self.bf9.exec(&mut InPlaceStore::new(&mut c_buffer));
        self.bf9.exec(&mut InPlaceStore::new(&mut s_buffer));

        let mut k = 0usize;
        let mut uk = 0usize;

        let q_modules = 9;
        let s = 2 * 27 / 3;

        {
            const S: usize = 8;
            let c_v = AvxStoreF::load(unsafe { c_buffer.get_unchecked(k..) });
            let s_v =
                AvxStoreF::load(unsafe { s_buffer.get_unchecked(q_modules - k - S..) }).reverse();
            let a_v = AvxStoreF::load(unsafe { a_buffer.get_unchecked(k..) });

            let twiddle_re = unsafe { *self.twiddle.get_unchecked(uk) };
            let twiddle_im = unsafe { *self.twiddle.get_unchecked(uk + 1) };

            let mut u0 = fma(c_v, twiddle_re, s_v * twiddle_im);
            let mut u1 = u0;
            let mut v0 = fma(c_v, twiddle_im, -s_v * twiddle_re);

            u0 += a_v;
            u1 *= f32::HALF;
            v0 *= f32::SQRT_3_OVER_2;
            u1 = u1 - a_v;

            let uc0 = u1 - v0;
            let uc1 = u1 + v0;

            u0.write(data.slice_from_mut(k..));
            uc1.write(data.slice_from_mut(s + k..));
            uc0.reverse().write(data.slice_from_mut(s - S - k..));

            k += 8;
            uk += 2;
        }

        {
            const S: usize = 1;
            let c_v = AvxStoreF::load1(unsafe { c_buffer.get_unchecked(k..) });
            let s_v = AvxStoreF::load1(unsafe { s_buffer.get_unchecked(q_modules - k - S..) });
            let a_v = AvxStoreF::load1(unsafe { a_buffer.get_unchecked(k..) });

            let twiddle_re = unsafe { *self.twiddle.get_unchecked(uk) };
            let twiddle_im = unsafe { *self.twiddle.get_unchecked(uk + 1) };

            let mut u0 = fma(c_v, twiddle_re, s_v * twiddle_im);
            let mut u1 = u0;
            let mut v0 = fma(c_v, twiddle_im, -s_v * twiddle_re);

            u0 += a_v;
            u1 *= f32::HALF;
            v0 *= f32::SQRT_3_OVER_2;
            u1 = u1 - a_v;

            let uc0 = u1 - v0;
            let uc1 = u1 + v0;

            u0.write1(data.slice_from_mut(k..));
            uc1.write1(data.slice_from_mut(s + k..));
            uc0.write1(data.slice_from_mut(s - S - k..));
        }
    }
}

impl AvxDct4Butterfly27f {
    #[target_feature(enable = "avx2", enable = "fma")]
    fn execute_impl(&self, data: &mut [f32]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(27) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }
        for chunk in data.chunks_exact_mut(27) {
            self.exec(&mut InPlaceStore::new(chunk));
        }
        Ok(())
    }

    #[target_feature(enable = "avx2", enable = "fma")]
    fn execute_into_impl(&self, input: &[f32], output: &mut [f32]) -> Result<(), PxdctError> {
        use crate::util::validate_oof_sizes;
        validate_oof_sizes!(input, output, 27);
        use crate::bidirectional::BiStore;
        for (src, dst) in input.chunks_exact(27).zip(output.chunks_exact_mut(27)) {
            self.exec(&mut BiStore::new(src, dst));
        }
        Ok(())
    }
}

impl PxdctExecutor<f32> for AvxDct4Butterfly27f {
    fn execute(&self, data: &mut [f32]) -> Result<(), PxdctError> {
        unsafe { self.execute_impl(data) }
    }

    fn execute_with_scratch(&self, data: &mut [f32], _: &mut [f32]) -> Result<(), PxdctError> {
        unsafe { self.execute_impl(data) }
    }

    fn execute_into(&self, input: &[f32], output: &mut [f32]) -> Result<(), PxdctError> {
        unsafe { self.execute_into_impl(input, output) }
    }

    fn execute_into_with_scratch(
        &self,
        input: &[f32],
        output: &mut [f32],
        _: &mut [f32],
    ) -> Result<(), PxdctError> {
        unsafe { self.execute_into_impl(input, output) }
    }

    fn length(&self) -> usize {
        27
    }

    fn scratch_size(&self) -> usize {
        0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::PxdctExecutor;
    use crate::avx::dct2_bf_power2::gen_test_avx_butterfly;
    use crate::tests::{naive_dct4, naive_dct4_f32};

    gen_test_avx_butterfly!(test_bf_dct4_3, AvxDct4Butterfly3, 3, 1e-7, naive_dct4);
    gen_test_avx_butterfly!(test_bf_dct4_9, AvxDct4Butterfly9, 9, 1e-7, naive_dct4);
    gen_test_avx_butterfly!(
        test_bf_dct4_27,
        AvxDct4Butterfly27f,
        27,
        1e-3,
        naive_dct4_f32
    );
}