lightweight-simd 0.1.0

// Most or all contents generated by AI (model glm-5.3)

//! Correctness tests for constructors, access, arithmetic operators, common
//! traits, and memory operations, checked against scalar-loop references.

use core::ops::Neg;
use num_traits::{Num, NumAssignOps, One, Zero};

use lightweight_simd::Simd;

macro_rules! check_all_n {
    ($func:ident, $t:ty, $mk:expr) => {
        $func::<$t, 1>(&$mk);
        $func::<$t, 2>(&$mk);
        $func::<$t, 3>(&$mk);
        $func::<$t, 4>(&$mk);
        $func::<$t, 5>(&$mk);
        $func::<$t, 8>(&$mk);
        $func::<$t, 16>(&$mk);
        $func::<$t, 32>(&$mk);
    };
}

macro_rules! check_all_n2 {
    ($func:ident, $t:ty, $a:expr, $b:expr) => {
        $func::<$t, 1>(&$a, &$b);
        $func::<$t, 2>(&$a, &$b);
        $func::<$t, 3>(&$a, &$b);
        $func::<$t, 4>(&$a, &$b);
        $func::<$t, 5>(&$a, &$b);
        $func::<$t, 8>(&$a, &$b);
        $func::<$t, 16>(&$a, &$b);
        $func::<$t, 32>(&$a, &$b);
    };
}

fn check_constructors<T, const N: usize>(mk: &impl Fn(usize) -> T)
where
    T: Num + Copy + PartialEq + std::fmt::Debug + Zero + One,
{
    let expect: [T; N] = core::array::from_fn(mk);
    let v = Simd::<T, N>::from_fn(mk);
    assert_eq!(v.to_array(), expect);
    assert_eq!(Simd::<T, N>::from(expect), v);
    assert_eq!(v.as_array(), &expect);
    assert_eq!(v.map(core::convert::identity).to_array(), expect);
    assert_eq!(Simd::<T, N>::splat(mk(3)).to_array(), core::array::from_fn(|_| mk(3)));
    let mut s = Simd::<T, N>::splat(mk(0));
    s.fill(mk(4));
    assert_eq!(s[0], mk(4));
    assert_eq!(s[N - 1], mk(4));
    let mut v2 = v;
    v2[1 % N] = mk(9);
    assert_eq!(v2[1 % N], mk(9));
    assert_eq!(Simd::<T, N>::zero().to_array(), core::array::from_fn(|_| T::zero()));
    assert_eq!(Simd::<T, N>::one()[0], T::one());
    assert!(<Simd<T, N> as Zero>::zero().is_zero());
    assert_eq!(<Simd<T, N> as One>::one()[0], T::one());
}

fn check_arith<T, const N: usize>(a_mk: &impl Fn(usize) -> T, b_mk: &impl Fn(usize) -> T)
where
    T: Num + NumAssignOps + Neg<Output = T> + Copy + PartialEq + std::fmt::Debug,
{
    let a = Simd::<T, N>::from_fn(a_mk);
    let b = Simd::<T, N>::from_fn(b_mk);

    // vector-vector binary
    assert_eq!((a + b).to_array(), core::array::from_fn(|i| a_mk(i) + b_mk(i)));
    assert_eq!((a - b).to_array(), core::array::from_fn(|i| a_mk(i) - b_mk(i)));
    assert_eq!((a * b).to_array(), core::array::from_fn(|i| a_mk(i) * b_mk(i)));
    assert_eq!((a / b).to_array(), core::array::from_fn(|i| a_mk(i) / b_mk(i)));

    // vector-scalar binary
    assert_eq!((a + b_mk(1)).to_array(), core::array::from_fn(|i| a_mk(i) + b_mk(1)));
    assert_eq!((a - b_mk(1)).to_array(), core::array::from_fn(|i| a_mk(i) - b_mk(1)));
    assert_eq!((a * b_mk(1)).to_array(), core::array::from_fn(|i| a_mk(i) * b_mk(1)));
    assert_eq!((a / b_mk(1)).to_array(), core::array::from_fn(|i| a_mk(i) / b_mk(1)));

    // vector-vector assign
    let mut c = a;
    c += b;
    assert_eq!(c.to_array(), core::array::from_fn(|i| a_mk(i) + b_mk(i)));
    c -= b;
    assert_eq!(c.to_array(), core::array::from_fn(|i| (a_mk(i) + b_mk(i)) - b_mk(i)));
    c *= b;
    assert_eq!(c.to_array(), core::array::from_fn(|i| ((a_mk(i) + b_mk(i)) - b_mk(i)) * b_mk(i)));
    c /= b;
    assert_eq!(c.to_array(), core::array::from_fn(|i| (((a_mk(i) + b_mk(i)) - b_mk(i)) * b_mk(i)) / b_mk(i)));

    // vector-scalar assign
    let mut c = a;
    c += b_mk(2);
    c -= b_mk(2);
    c *= b_mk(2);
    c /= b_mk(2);
    assert_eq!(c.to_array(), core::array::from_fn(|i| (((a_mk(i) + b_mk(2)) - b_mk(2)) * b_mk(2)) / b_mk(2)));

    // negation
    assert_eq!((-a).to_array(), core::array::from_fn(|i| -a_mk(i)));

    // num trait pass-through
    assert!(Simd::<T, N>::zero().is_zero());
}

fn check_memory<T, const N: usize>(mk: &impl Fn(usize) -> T)
where
    T: Zero + Copy + PartialEq + std::fmt::Debug,
{
    let full: Vec<T> = (0..N).map(mk).collect();
    let v = Simd::<T, N>::from_slice(&full);
    assert_eq!(v.to_array(), core::array::from_fn(mk));

    // padded loads
    for k in 0..=N {
        let pv = Simd::<T, N>::from_slice_pad(&full[..k]);
        let expect: [T; N] = core::array::from_fn(|i| if i < k { mk(i) } else { T::zero() });
        assert_eq!(pv.to_array(), expect);
    }

    // full store
    let mut buf = vec![T::zero(); N];
    v.store_slice(&mut buf);
    assert_eq!(buf, full);

    // partial stores
    for k in 0..=(N + 2) {
        let mut buf = vec![T::zero(); k];
        v.store_slice_partial(&mut buf);
        let n_valid = k.min(N);
        assert_eq!(&buf[..n_valid], &full[..n_valid]);
        for x in &buf[n_valid..] {
            assert!(x.is_zero());
        }
    }
}

#[test]
fn constructors_f64() {
    check_all_n!(check_constructors, f64, |i| i as f64 * 0.5 + 1.5);
}

#[test]
fn constructors_other_types() {
    check_constructors::<f32, 4>(&|i| i as f32 * 0.25 - 1.0);
    check_constructors::<i64, 3>(&|i| i as i64 * 7 - 3);
    check_constructors::<i32, 5>(&|i| i as i32 + 9);
}

#[test]
fn constructors_complex_and_half() {
    use half::f16;
    use num_complex::Complex;
    check_constructors::<Complex<f64>, 4>(&|i| Complex::new(i as f64, 1.5));
    check_constructors::<f16, 4>(&|i| f16::from_f32(i as f32 * 0.5 + 1.0));
}

#[test]
fn arith_f64() {
    check_all_n2!(check_arith, f64, |i| i as f64 * 0.5 - 1.0, |i| i as f64 * 0.25 + 2.0);
}

#[test]
fn arith_other_types() {
    check_all_n2!(check_arith, i64, |i| i as i64 * 3 - 7, |i| i as i64 % 5 + 7);
    check_all_n2!(check_arith, i32, |i| i as i32 * 3 - 7, |i| i as i32 % 5 + 7);
    check_arith::<f32, 4>(&|i| i as f32 * 0.5 - 1.0, &|i| i as f32 * 0.25 + 2.0);
    check_arith::<f32, 3>(&|i| i as f32 * 0.5 - 1.0, &|i| i as f32 * 0.25 + 2.0);
}

#[test]
fn arith_complex_and_half() {
    use half::{bf16, f16};
    use num_complex::Complex;
    check_arith::<Complex<f64>, 4>(&|i| Complex::new(i as f64, -(i as f64) * 2.0), &|i| {
        Complex::new(i as f64 + 1.5, 0.25)
    });
    check_arith::<Complex<f64>, 3>(&|i| Complex::new(i as f64, -(i as f64) * 2.0), &|i| {
        Complex::new(i as f64 + 1.5, 0.25)
    });
    check_arith::<f16, 4>(&|i| f16::from_f32(i as f32 * 0.5 - 1.0), &|i| f16::from_f32(i as f32 * 0.25 + 2.0));
    check_arith::<bf16, 5>(&|i| bf16::from_f32(i as f32 * 0.5 - 1.0), &|i| bf16::from_f32(i as f32 * 0.25 + 2.0));
}

#[test]
fn memory_f64() {
    check_all_n!(check_memory, f64, |i| i as f64 * 0.5 + 1.5);
}

#[test]
fn memory_other_types() {
    check_memory::<f32, 8>(&|i| i as f32 * 0.25);
    check_memory::<i32, 5>(&|i| i as i32 + 3);
}

#[test]
fn common_traits_i64() {
    use std::collections::HashSet;
    let a = Simd::<i64, 4>::from([1, 2, 3, 4]);
    let b = Simd::<i64, 4>::from([1, 2, 3, 4]);
    let c = Simd::<i64, 4>::from([1, 2, 3, 5]);
    assert_eq!(a, b);
    assert_ne!(a, c);
    assert_eq!(Simd::<i64, 4>::default(), Simd::<i64, 4>::zero());
    let mut set = HashSet::new();
    set.insert(a);
    assert!(set.contains(&b));
    assert!(!set.contains(&c));
    assert!(format!("{a:?}").starts_with("Simd(["));
}

#[test]
fn uninit_scratch() {
    let mut v = unsafe { Simd::<f64, 4>::uninit() };
    v.fill(3.0);
    assert_eq!(v.reduce_sum(), 12.0);
}

#[test]
#[should_panic]
fn from_slice_too_short() {
    let _ = Simd::<f64, 4>::from_slice(&[0.0; 3]);
}

#[test]
#[should_panic]
fn store_slice_too_short() {
    Simd::<f64, 4>::zero().store_slice(&mut [0.0; 3]);
}

#[test]
#[should_panic]
fn reduce_sum_n_out_of_range() {
    Simd::<f64, 4>::zero().reduce_sum_n(5);
}