fidget-core 0.5.1

//! Test suites for each evaluator type
pub mod float_slice;
pub mod grad_slice;
pub mod interval;
pub mod point;

use crate::types::FloatExt;

// Internal-only tests
#[cfg(test)] // not enabled for eval-tests
mod symbolic_deriv;

use crate::{
    context::{Context, IntoNode, Node},
    eval::Tape,
    var::Var,
};

/// Builds a function which stresses the register allocator and function caller
pub fn build_stress_fn(n: usize) -> (Context, Node) {
    let mut inputs = vec![];
    let mut ctx = Context::new();
    let mut sum = ctx.constant(0.0);
    let x = ctx.x();
    let y = ctx.y();
    let z = ctx.z();

    // Build up the sum (x * 1) + (y * 2) + (z * 3) + (x * 4) + ...
    //
    // We're going to both send these operations into a sin(..) node, then add
    // them afterwards (in reverse order), meaning their allocations must
    // persist beyond the sine.
    for i in 1..=n {
        let d = ctx.mul(i as f32, [x, y, z][i % 3]).unwrap();
        inputs.push(d);
        sum = ctx.add(sum, d).unwrap();
    }

    sum = ctx.sin(sum).unwrap();
    for i in inputs.into_iter().rev() {
        sum = ctx.add(sum, i).unwrap();
    }

    (ctx, sum)
}

/// Pick a bunch of arguments, some of which are spicy
fn test_args_n(n: i64) -> Vec<f32> {
    let mut args = (-n..=n)
        .map(|i| std::f32::consts::PI * 2.0 * i as f32 / (n as f32))
        .collect::<Vec<_>>();
    args.push(1.0);
    args.push(5.0);
    args.push(0.5);
    args.push(1.5);
    args.push(10.0);
    args.push(std::f32::consts::PI);
    args.push(std::f32::consts::FRAC_PI_2);
    args.push(std::f32::consts::FRAC_1_PI);
    args.push(std::f32::consts::SQRT_2);
    args.push(f32::NAN);
    args
}

fn test_args() -> Vec<f32> {
    test_args_n(32)
}

fn bind_xy<T: Tape, V, G: Into<V>>(tape: &T) -> Box<dyn Fn(G, G) -> [V; 2]> {
    let vars = tape.vars();
    let ix = vars[&Var::X];
    let iy = vars[&Var::Y];
    assert_ne!(ix, iy);
    Box::new(move |x, y| {
        let mut out = [None, None];
        out[ix] = Some(x.into());
        out[iy] = Some(y.into());
        out.map(Option::unwrap)
    })
}

fn bind_xyz<T: Tape, V, G: Into<V>>(
    tape: &T,
) -> Box<dyn Fn(G, G, G) -> [V; 3]> {
    let vars = tape.vars();
    let ix = vars[&Var::X];
    let iy = vars[&Var::Y];
    let iz = vars[&Var::Z];
    assert_ne!(ix, iy);
    assert_ne!(iy, iz);
    assert_ne!(ix, iz);
    Box::new(move |x, y, z| {
        let mut out = [None, None, None];
        out[ix] = Some(x.into());
        out[iy] = Some(y.into());
        out[iz] = Some(z.into());
        out.map(Option::unwrap)
    })
}

/// Trait for canonical evaluation testing of unary operations
pub trait CanonicalUnaryOp {
    const NAME: &'static str;
    fn build(ctx: &mut Context, arg: Node) -> Node;
    fn eval_f32(arg: f32) -> f32;

    /// Returns true if there is a bidirectional discontinuity at a position
    ///
    /// This means that we should skip gradient checking, because we can't
    /// accurately estimate the gradient on either side.
    fn discontinuous_at(_lhs: f32) -> bool {
        false
    }
}

/// Trait for canonical evaluation testing of binary operations
pub trait CanonicalBinaryOp {
    const NAME: &'static str;
    fn build<A: IntoNode, B: IntoNode>(
        ctx: &mut Context,
        lhs: A,
        rhs: B,
    ) -> Node;
    fn eval(lhs: f32, rhs: f32) -> f32;

    /// Returns true if there is a bidirectional discontinuity at a position
    ///
    /// This means that we should skip gradient checking, because we can't
    /// accurately estimate the gradient on either side.
    fn discontinuous_at(_lhs: f32, _rhs: f32) -> bool {
        false
    }
}

macro_rules! declare_canonical_unary {
    (Context::$i:ident, |$a:ident| $t:expr, |$b:ident| $u:expr) => {
        pub struct $i;
        impl CanonicalUnaryOp for $i {
            const NAME: &'static str = stringify!($i);
            fn build(ctx: &mut Context, arg: Node) -> Node {
                Context::$i(ctx, arg).unwrap()
            }
            fn eval_f32($a: f32) -> f32 {
                $t
            }
            fn discontinuous_at($b: f32) -> bool {
                $u
            }
        }
    };
    (Context::$i:ident, |$lhs:ident| $t:expr) => {
        declare_canonical_unary!(Context::$i, |$lhs| $t, |_a| false);
    };
}

macro_rules! declare_canonical_binary {
    (Context::$i:ident, |$lhs:ident, $rhs:ident| $t:expr, |$lhs2:ident, $rhs2: ident| $d:expr) => {
        pub struct $i;
        impl CanonicalBinaryOp for $i {
            const NAME: &'static str = stringify!($i);
            fn build<A: IntoNode, B: IntoNode>(
                ctx: &mut Context,
                lhs: A,
                rhs: B,
            ) -> Node {
                let lhs = lhs.into_node(ctx).unwrap();
                let rhs = rhs.into_node(ctx).unwrap();
                Context::$i(ctx, lhs, rhs).unwrap()
            }
            fn eval($lhs: f32, $rhs: f32) -> f32 {
                $t
            }
            fn discontinuous_at($lhs2: f32, $rhs2: f32) -> bool {
                $d
            }
        }
    };
    (Context::$i:ident, |$lhs:ident, $rhs:ident| $t:expr) => {
        declare_canonical_binary!(Context::$i, |$lhs, $rhs| $t, |_a, _b| false);
    };
}

#[allow(
    non_camel_case_types,
    clippy::useless_conversion,
    clippy::unnecessary_cast
)]
pub mod canonical {
    use super::*;

    declare_canonical_unary!(Context::neg, |a| -a);
    declare_canonical_unary!(Context::recip, |a| 1.0 / a);
    declare_canonical_unary!(Context::abs, |a| a.abs());
    declare_canonical_unary!(Context::sin, |a| a.sin());
    declare_canonical_unary!(Context::cos, |a| a.cos());
    declare_canonical_unary!(Context::tan, |a| a.tan());
    declare_canonical_unary!(Context::asin, |a| a.asin());
    declare_canonical_unary!(Context::acos, |a| a.acos());
    declare_canonical_unary!(Context::atan, |a| a.atan());
    declare_canonical_unary!(Context::exp, |a| a.exp());
    declare_canonical_unary!(Context::ln, |a| a.ln());
    declare_canonical_unary!(Context::square, |a| a * a);
    declare_canonical_unary!(Context::sqrt, |a| a.sqrt());
    declare_canonical_unary!(Context::floor, |a| a.floor());
    declare_canonical_unary!(Context::ceil, |a| a.ceil());
    declare_canonical_unary!(Context::round, |a| a.round());
    declare_canonical_unary!(Context::not, |a| a.not(), |a| a == 0.0);
    declare_canonical_unary!(
        Context::rand,
        |a| a.rand(),
        |_a| true // always discontinuous, so ignore all gradients
    );

    declare_canonical_binary!(Context::add, |a, b| a + b);
    declare_canonical_binary!(Context::sub, |a, b| a - b);
    declare_canonical_binary!(Context::mul, |a, b| a * b);
    declare_canonical_binary!(Context::div, |a, b| a / b);
    declare_canonical_binary!(Context::min, |a, b| a.min_choice(b).0);
    declare_canonical_binary!(Context::max, |a, b| a.max_choice(b).0);
    declare_canonical_binary!(Context::compare, |a, b| a.compare(b), |a, b| a
        == b);
    declare_canonical_binary!(
        Context::modulo,
        |a, b| a.rem_euclid(b),
        |a, b| {
            let v = a / b;
            (v.round() - v).abs() < 1e-9
        }
    );
    declare_canonical_binary!(
        Context::and,
        |a, b| a.and_choice(b).0,
        |a, _b| a == 0.0 // discontinuity, because either side snaps to b
    );
    declare_canonical_binary!(
        Context::or,
        |a, b| a.or_choice(b).0,
        |a, _b| a == 0.0 // discontinuity, because either side snaps to a
    );
    declare_canonical_binary!(Context::atan2, |y, x| y.atan2(x));
    declare_canonical_binary!(
        Context::mix,
        |a, b| a.mix(b),
        |_a, _b| true // always discontinuous
    );
}

#[macro_export]
macro_rules! one_unary_test {
    ($tester:ty, $i:ident) => {
        #[test]
        fn $i() {
            <$tester>::test_unary::<$crate::eval::test::canonical::$i>()
        }
    };
}

#[macro_export]
macro_rules! one_binary_test {
    ($tester:ty, $i:ident) => {
        #[test]
        fn $i() {
            <$tester>::test_binary::<$crate::eval::test::canonical::$i>()
        }
    };
}

#[macro_export]
macro_rules! all_unary_tests {
    ($tester:ty) => {
        $crate::one_unary_test!($tester, neg);
        $crate::one_unary_test!($tester, recip);
        $crate::one_unary_test!($tester, abs);
        $crate::one_unary_test!($tester, sin);
        $crate::one_unary_test!($tester, cos);
        $crate::one_unary_test!($tester, tan);
        $crate::one_unary_test!($tester, asin);
        $crate::one_unary_test!($tester, acos);
        $crate::one_unary_test!($tester, atan);
        $crate::one_unary_test!($tester, exp);
        $crate::one_unary_test!($tester, ln);
        $crate::one_unary_test!($tester, not);
        $crate::one_unary_test!($tester, square);
        $crate::one_unary_test!($tester, floor);
        $crate::one_unary_test!($tester, ceil);
        $crate::one_unary_test!($tester, round);
        $crate::one_unary_test!($tester, sqrt);
        $crate::one_unary_test!($tester, rand);
    };
}

#[macro_export]
macro_rules! all_binary_tests {
    ($tester:ty) => {
        $crate::one_binary_test!($tester, add);
        $crate::one_binary_test!($tester, sub);
        $crate::one_binary_test!($tester, mul);
        $crate::one_binary_test!($tester, div);
        $crate::one_binary_test!($tester, atan2);
        $crate::one_binary_test!($tester, min);
        $crate::one_binary_test!($tester, max);
        $crate::one_binary_test!($tester, compare);
        $crate::one_binary_test!($tester, modulo);
        $crate::one_binary_test!($tester, and);
        $crate::one_binary_test!($tester, or);
        $crate::one_binary_test!($tester, mix);
    };
}