macerator 0.5.1

use std::{vec, vec::Vec};

use half::f16;
use num_traits::{Float, NumCast};
use rand::distr::uniform::SampleUniform;

use crate::{
    tests::{binop, random, test_binop, SIZE},
    vload_unaligned, Mask, Scalar, Simd, VEq, VOrd, Vector,
};

pub(crate) trait CmpOp<T: Scalar> {
    fn call<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Mask<S, T>;
}

#[inline(always)]
fn test_cmp<S: Simd, T: Scalar, Op: CmpOp<T>>(lhs: &[T], rhs: &[T]) -> Vec<bool> {
    let lanes = T::lanes::<S>();
    let mut output = vec![false; lhs.len()];
    let lhs = lhs.chunks_exact(lanes);
    let rhs = rhs.chunks_exact(lanes);
    let out = output.chunks_exact_mut(lanes);
    for ((lhs, rhs), out) in lhs.zip(rhs).zip(out) {
        let lhs = unsafe { vload_unaligned(lhs.as_ptr()) };
        let rhs = unsafe { vload_unaligned(rhs.as_ptr()) };

        unsafe { T::mask_store_as_bool(out.as_mut_ptr(), Op::call::<S>(lhs, rhs)) };
    }
    output
}

macro_rules! cmp_op {
    ($trait: ident, $impl: expr) => {
        ::paste::paste! {
            struct [<$trait Op>]<T>(::core::marker::PhantomData<T>);
            impl<T: $trait> CmpOp<T> for [<$trait Op>]<T> {
                #[inline(always)]
                fn call<S: Simd>(lhs: $crate::Vector<S, T>, rhs: $crate::Vector<S, T>) -> Mask<S, T> {
                    $impl(lhs, rhs)
                }
            }
        }
    };
}

/// Random operands in `0..127`.
fn random_inputs<T: SampleUniform + NumCast + Copy>() -> (Vec<T>, Vec<T>) {
    let (lo, hi) = (T::from(0).unwrap(), T::from(127).unwrap());
    (random(lo, hi), random(lo, hi))
}

/// Every pairing of NaN (both signs), signed zero, signed infinity and ordinary
/// values, which the random operands never produce.
fn special_float_inputs<T: Float>() -> (Vec<T>, Vec<T>) {
    let values = [
        T::nan(),
        T::one(),
        -T::one(),
        T::zero(),
        -T::zero(),
        T::infinity(),
        T::neg_infinity(),
        -T::nan(),
    ];
    let n = values.len();
    let lhs = (0..SIZE).map(|i| values[i % n]).collect();
    let rhs = (0..SIZE).map(|i| values[(i / n) % n]).collect();
    (lhs, rhs)
}

macro_rules! testgen_cmp {
    ($test_fn: ident, $reference: expr, $($(#[$meta:meta])* $ty: ty),*) => {
        testgen_cmp!(@inputs random_inputs; $test_fn, $reference, $($(#[$meta])* $ty),*);
    };
    (@inputs $inputs: path; $test_fn: ident, $reference: expr, $($(#[$meta:meta])* $ty: ty),*) => {
        $(::paste::paste! {
            $(#[$meta])*
            #[::wasm_bindgen_test::wasm_bindgen_test(unsupported = test)]
            fn [<$test_fn _ $ty>]() {
                let (lhs, rhs): (Vec<$ty>, Vec<$ty>) = $inputs();
                let out_ref = lhs
                    .iter()
                    .zip(rhs.iter())
                    .map(|(a, b)| $ty::$reference(a, b))
                    .collect::<Vec<_>>();
                #[cfg(x86)]
                unsafe {
                    use $crate::backend::x86::*;
                    #[cfg(avx512_fp16)]
                    if V4FP16::is_available() {
                        let out = V4FP16::run_vectorized(|| [<$test_fn _impl>]::<V4FP16, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    #[cfg(avx512)]
                    if V4::is_available() {
                        let out = V4::run_vectorized(|| [<$test_fn _impl>]::<V4, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    if V3::is_available() {
                        let out = V3::run_vectorized(|| [<$test_fn _impl>]::<V3, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    if V2::is_available() {
                        let out = V2::run_vectorized(|| [<$test_fn _impl>]::<V2, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                }
                #[cfg(aarch64)]
                unsafe {
                    use $crate::backend::aarch64::*;
                    #[cfg(feature = "fp16")]
                    if NeonFP16::is_available() {
                        let out = NeonFP16::run_vectorized(|| [<$test_fn _impl>]::<NeonFP16, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    if NeonFma::is_available() {
                        let out = NeonFma::run_vectorized(|| [<$test_fn _impl>]::<NeonFma, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                }
                #[cfg(loong64)]
                unsafe {
                    use $crate::backend::loong64::*;
                    if Lasx::is_available() {
                        let out = Lasx::run_vectorized(|| [<$test_fn _impl>]::<Lasx, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                    if Lsx::is_available() {
                        let out = Lsx::run_vectorized(|| [<$test_fn _impl>]::<Lsx, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                }
                #[cfg(wasm32)]
                unsafe {
                    use crate::backend::wasm32;
                    #[cfg(relaxed_simd)]
                    if wasm32::Simd128Relaxed::is_available() {
                        let out = wasm32::Simd128Relaxed::run_vectorized(|| [<$test_fn _impl>]::<wasm32::Simd128Relaxed, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                    if wasm32::Simd128Fallback::is_available() {
                        let out = wasm32::Simd128Fallback::run_vectorized(|| [<$test_fn _impl>]::<wasm32::Simd128Fallback, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                }
                let out = [<$test_fn _impl>]::<$crate::backend::scalar::Fallback, $ty>(&lhs, &rhs);
                assert_eq!(out_ref, out);
            }
        })*
    };
}

macro_rules! testgen_min_max {
    ($test_fn: ident, $reference: expr, $($(#[$meta:meta])* $ty: ty),*) => {
        $(::paste::paste! {
            $(#[$meta])*
            #[::wasm_bindgen_test::wasm_bindgen_test(unsupported = test)]
            fn [<$test_fn _ $ty>]() {
                let (lhs, rhs): (Vec<$ty>, Vec<$ty>) = random_inputs();
                let out_ref = lhs
                    .iter()
                    .zip(rhs.iter())
                    .map(|(a, b)| $ty::$reference(*a, *b))
                    .collect::<Vec<_>>();
                #[cfg(x86)]
                unsafe {
                    use $crate::backend::x86::*;
                    #[cfg(avx512_fp16)]
                    if V4FP16::is_available() {
                        let out = V4FP16::run_vectorized(|| [<$test_fn _impl>]::<V4FP16, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    #[cfg(avx512)]
                    if V4::is_available() {
                        let out = V4::run_vectorized(|| [<$test_fn _impl>]::<V4, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    if V3::is_available() {
                        let out = V3::run_vectorized(|| [<$test_fn _impl>]::<V3, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    if V2::is_available() {
                        let out = V2::run_vectorized(|| [<$test_fn _impl>]::<V2, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                }
                #[cfg(aarch64)]
                unsafe {
                    use $crate::backend::aarch64::*;
                    #[cfg(feature = "fp16")]
                    if NeonFP16::is_available() {
                        let out = NeonFP16::run_vectorized(|| [<$test_fn _impl>]::<NeonFP16, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                    if NeonFma::is_available() {
                        let out = NeonFma::run_vectorized(|| [<$test_fn _impl>]::<NeonFma, $ty>(&lhs, &rhs));
                        assert_eq!(out_ref, out);
                    }
                }
                #[cfg(loong64)]
                unsafe {
                    use $crate::backend::loong64::*;
                    if Lasx::is_available() {
                        let out = Lasx::run_vectorized(|| [<$test_fn _impl>]::<Lasx, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                    if Lsx::is_available() {
                        let out = Lsx::run_vectorized(|| [<$test_fn _impl>]::<Lsx, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                }
                #[cfg(wasm32)]
                unsafe {
                    use crate::backend::wasm32;
                    #[cfg(relaxed_simd)]
                    if wasm32::Simd128Relaxed::is_available() {
                        let out = wasm32::Simd128Relaxed::run_vectorized(|| [<$test_fn _impl>]::<wasm32::Simd128Relaxed, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                    if wasm32::Simd128Fallback::is_available() {
                        let out = wasm32::Simd128Fallback::run_vectorized(|| [<$test_fn _impl>]::<wasm32::Simd128Fallback, $ty>(&lhs, &rhs));
                        assert_eq!(&out_ref, &out);
                    }
                }
                let out = [<$test_fn _impl>]::<$crate::backend::scalar::Fallback, $ty>(&lhs, &rhs);
                assert_eq!(out_ref, out);
            }
        })*
    };
}

#[inline(always)]
fn test_eq_impl<S: Simd, T: VEq>(lhs: &[T], rhs: &[T]) -> Vec<bool> {
    cmp_op!(VEq, |a: Vector<S, T>, b| a.eq(b));
    test_cmp::<S, T, VEqOp<T>>(lhs, rhs)
}

#[inline(always)]
fn test_ne_impl<S: Simd, T: VEq>(lhs: &[T], rhs: &[T]) -> Vec<bool> {
    cmp_op!(VEq, |a: Vector<S, T>, b| a.ne(b));
    test_cmp::<S, T, VEqOp<T>>(lhs, rhs)
}

#[inline(always)]
fn test_lt_impl<S: Simd, T: VOrd>(lhs: &[T], rhs: &[T]) -> Vec<bool> {
    cmp_op!(VOrd, |a: Vector<S, T>, b| a.lt(b));
    test_cmp::<S, T, VOrdOp<T>>(lhs, rhs)
}

#[inline(always)]
fn test_le_impl<S: Simd, T: VOrd>(lhs: &[T], rhs: &[T]) -> Vec<bool> {
    cmp_op!(VOrd, |a: Vector<S, T>, b| a.le(b));
    test_cmp::<S, T, VOrdOp<T>>(lhs, rhs)
}

#[inline(always)]
fn test_gt_impl<S: Simd, T: VOrd>(lhs: &[T], rhs: &[T]) -> Vec<bool> {
    cmp_op!(VOrd, |a: Vector<S, T>, b| a.gt(b));
    test_cmp::<S, T, VOrdOp<T>>(lhs, rhs)
}

#[inline(always)]
fn test_ge_impl<S: Simd, T: VOrd>(lhs: &[T], rhs: &[T]) -> Vec<bool> {
    cmp_op!(VOrd, |a: Vector<S, T>, b| a.ge(b));
    test_cmp::<S, T, VOrdOp<T>>(lhs, rhs)
}

#[inline(always)]
fn test_min_impl<S: Simd, T: VOrd>(lhs: &[T], rhs: &[T]) -> Vec<T> {
    binop!(VOrd, |a: Vector<S, T>, b| a.min(b));
    test_binop::<S, T, VOrdOp<T>>(lhs, rhs)
}

#[inline(always)]
fn test_max_impl<S: Simd, T: VOrd>(lhs: &[T], rhs: &[T]) -> Vec<T> {
    binop!(VOrd, |a: Vector<S, T>, b| a.max(b));
    test_binop::<S, T, VOrdOp<T>>(lhs, rhs)
}

testgen_cmp!(
    test_eq,
    eq,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f64
);
testgen_cmp!(
    test_ne,
    ne,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f64
);
testgen_cmp!(
    test_lt,
    lt,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f64
);
testgen_cmp!(
    test_gt,
    gt,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f64
);
testgen_cmp!(
    test_le,
    le,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f64
);
testgen_cmp!(
    test_ge,
    ge,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
    f64
);
testgen_min_max!(
    test_min,
    min,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(aarch64, x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(aarch64, x86_v4, loong64)), ignore)]
    f64
);
testgen_min_max!(
    test_max,
    max,
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    #[cfg_attr(all(miri, any(aarch64, x86_v4, loong64)), ignore)]
    f32,
    u64,
    i64,
    #[cfg_attr(all(miri, any(aarch64, x86_v4, loong64)), ignore)]
    f64
);

/// The tests above draw operands from `0..127`, so they never see negative
/// signed values or unsigned values with the top bit set, which is where the
/// biased compares emulating unsigned ordering on x86 and wasm matter.
mod full_range {
    use super::*;
    use num_traits::Bounded;

    /// Operands over the whole range of the type, starting with every pairing
    /// of `MIN`, `MAX`, 0 and 1.
    fn random_inputs<T: SampleUniform + NumCast + Copy + Bounded>() -> (Vec<T>, Vec<T>) {
        let mut lhs = random(T::min_value(), T::max_value());
        let mut rhs = random(T::min_value(), T::max_value());
        let edges = [
            T::min_value(),
            T::max_value(),
            T::from(0).unwrap(),
            T::from(1).unwrap(),
        ];
        for (i, a) in edges.iter().enumerate() {
            for (j, b) in edges.iter().enumerate() {
                lhs[i * edges.len() + j] = *a;
                rhs[i * edges.len() + j] = *b;
            }
        }
        (lhs, rhs)
    }

    testgen_cmp!(test_eq, eq, u8, i8, u16, i16, u32, i32, u64, i64);
    testgen_cmp!(test_ne, ne, u8, i8, u16, i16, u32, i32, u64, i64);
    testgen_cmp!(test_lt, lt, u8, i8, u16, i16, u32, i32, u64, i64);
    testgen_cmp!(test_le, le, u8, i8, u16, i16, u32, i32, u64, i64);
    testgen_cmp!(test_gt, gt, u8, i8, u16, i16, u32, i32, u64, i64);
    testgen_cmp!(test_ge, ge, u8, i8, u16, i16, u32, i32, u64, i64);
    testgen_min_max!(test_min, min, u8, i8, u16, i16, u32, i32, u64, i64);
    testgen_min_max!(test_max, max, u8, i8, u16, i16, u32, i32, u64, i64);
}

mod nan {
    use super::*;

    testgen_cmp!(
        @inputs special_float_inputs;
        test_eq,
        eq,
        f16,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f32,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f64
    );
    testgen_cmp!(
        @inputs special_float_inputs;
        test_ne,
        ne,
        f16,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f32,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f64
    );
    testgen_cmp!(
        @inputs special_float_inputs;
        test_lt,
        lt,
        f16,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f32,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f64
    );
    testgen_cmp!(
        @inputs special_float_inputs;
        test_le,
        le,
        f16,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f32,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f64
    );
    testgen_cmp!(
        @inputs special_float_inputs;
        test_gt,
        gt,
        f16,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f32,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f64
    );
    testgen_cmp!(
        @inputs special_float_inputs;
        test_ge,
        ge,
        f16,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f32,
        #[cfg_attr(all(miri, any(x86_v4, loong64)), ignore)]
        f64
    );
}