exp_rs/

expression_functions.rs

//! Expression functions implementation for the exp-rs library.
//!
//! This module provides functionality for defining and evaluating functions
//! that are specified as expression strings rather than native Rust code.

extern crate alloc;
use crate::Real;
use crate::context::EvalContext;
use crate::error::Result;
use crate::eval::eval_ast;
use crate::types::AstExpr;
#[cfg(not(test))]
use alloc::rc::Rc;
#[cfg(test)]
use std::rc::Rc;
use alloc::borrow::Cow;
use alloc::string::ToString;

/// Evaluates an expression function with the given arguments.
///
/// This is a helper function used internally by the evaluation logic.
pub fn eval_expression_function<'a>(
    ast: &AstExpr,
    param_names: &[Cow<'a , str>],
    arg_values: &[Real],
    parent_ctx: Option<Rc<EvalContext<'a>>>,
) -> Result<Real> {
    let mut temp_ctx = EvalContext::new();
    if let Some(parent) = parent_ctx {
        temp_ctx.parent = Some(Rc::clone(&parent));
    }
    for (param_name, &arg_val) in param_names.iter().zip(arg_values.iter()) {
        temp_ctx.variables.insert(param_name.to_string(), arg_val);
    }
    eval_ast(ast, Some(Rc::new(temp_ctx)))
}


#[cfg(test)]
mod tests {
    use super::*;
    use crate::constants;
    use crate::engine::interp;
    // Import the macro into the test module scope
    use crate::assert_approx_eq;
    // Import Real for casting
    use crate::Real;

    #[test]
    fn test_simple_expression_function() {
        let mut ctx = EvalContext::new();

        ctx.register_expression_function("double", &["x"], "x * 2")
            .unwrap();

        let result = interp("double(5)", Some(Rc::new(ctx.clone()))).unwrap();
        assert_eq!(result, 10.0);
    }

    #[test]
    fn test_nested_expression_functions() {
        let mut ctx = EvalContext::new();

        ctx.register_expression_function("square", &["x"], "x * x")
            .unwrap();
        ctx.register_expression_function("cube", &["x"], "x * square(x)")
            .unwrap();

        println!("Registered expression functions: {:?}", ctx.function_registry.expression_functions.keys().collect::<Vec<_>>());

        assert!(ctx.function_registry.expression_functions.contains_key("square"), "Context missing 'square' function");
        assert!(ctx.function_registry.expression_functions.contains_key("cube"), "Context missing 'cube' function");

        let ast = crate::engine::parse_expression("cube(3)").unwrap();
        println!("Parsed expression: {:?}", ast);

        let result = interp("cube(3)", Some(Rc::new(ctx.clone())));
        match &result {
            Ok(val) => assert_eq!(*val, 27.0),
            Err(e) => {
                println!("Error evaluating cube(3): {:?}", e);
                println!("Context expression functions: {:?}", ctx.function_registry.expression_functions.keys().collect::<Vec<_>>());
                println!("Full context: variables={:?}, constants={:?}, expression_functions={:?}", ctx.variables, ctx.constants, ctx.function_registry.expression_functions.keys().collect::<Vec<_>>());
                panic!("Failed to evaluate cube(3): {:?}", e);
            }
        }
    }

    #[test]
    fn test_expression_function_with_multiple_params() {
        let mut ctx = EvalContext::new();

        ctx.register_expression_function("weighted_sum", &["a", "b", "w"], "a * w + b * (1 - w)")
            .unwrap();

        let body_ast = crate::engine::parse_expression_with_reserved(
            "a * w + b * (1 - w)",
            Some(&["a".to_string(), "b".to_string(), "w".to_string()])
        ).unwrap();
        println!("AST for function body 'a * w + b * (1 - w)': {:?}", body_ast);

        fn assert_no_function_w(ast: &AstExpr) {
            match ast {
                AstExpr::Function { name, args } => {
                    assert_ne!(name, "w", "Parameter 'w' should not be parsed as a function");
                    for arg in args {
                        assert_no_function_w(arg);
                    }
                }
                AstExpr::Array { index, .. } => assert_no_function_w(index),
                _ => {}
            }
        }
        assert_no_function_w(&body_ast);

        let w_ast = crate::engine::parse_expression_with_reserved(
            "w",
            Some(&["w".to_string()])
        ).unwrap();
        println!("AST for 'w': {:?}", w_ast);
        match w_ast {
            AstExpr::Variable(ref name) => assert_eq!(name, "w"),
            _ => panic!("Expected variable node for 'w'"),
        }

        let w_b_ast = crate::engine::parse_expression_with_reserved(
            "w b",
            Some(&["w".to_string()])
        );
        println!("AST for 'w b': {:?}", w_b_ast);
        assert!(w_b_ast.is_err(), "Expected parse error for 'w b' when 'w' is a reserved parameter");

        let ast = crate::engine::parse_expression("weighted_sum(10, 20, 0.3)");
        match ast {
            Ok(ast) => println!("Parsed expression: {:?}", ast),
            Err(e) => println!("Parse error for weighted_sum(10, 20, 0.3): {:?}", e),
        }

        let result1 = interp("weighted_sum(10, 20, 0.3)", Some(Rc::new(ctx.clone())));
        match result1 {
            Ok(val) => assert_eq!(val, 10.0 * 0.3 + 20.0 * 0.7),
            Err(e) => {
                println!("Error evaluating weighted_sum(10, 20, 0.3): {:?}", e);
                panic!("Failed to evaluate weighted_sum(10, 20, 0.3): {:?}", e);
            }
        }

        let result2 = interp("weighted_sum(10, 20, 0.7)", Some(Rc::new(ctx.clone())));
        match result2 {
            Ok(val) => assert_eq!(val, 10.0 * 0.7 + 20.0 * 0.3),
            Err(e) => {
                println!("Error evaluating weighted_sum(10, 20, 0.7): {:?}", e);
                panic!("Failed to evaluate weighted_sum(10, 20, 0.7): {:?}", e);
            }
        }
    }

    #[test]
    fn test_expression_function_with_context_variables() {
        let mut ctx = EvalContext::new();

        ctx.variables.insert("base".to_string().into(), 10.0);
        ctx.constants.insert("FACTOR".to_string().into(), 2.5);

        println!("Context variables before: {:?}", ctx.variables);
        println!("Context constants before: {:?}", ctx.constants);

        ctx.register_expression_function("scaled_value", &["x"], "base + x * FACTOR")
            .unwrap();

        println!("Context variables after: {:?}", ctx.variables);
        println!("Context constants after: {:?}", ctx.constants);

        assert!(ctx.variables.contains_key("base"), "Context missing 'base' variable");
        assert!(ctx.constants.contains_key("FACTOR"), "Context missing 'FACTOR' constant");

        let ast = crate::engine::parse_expression("scaled_value(4)").unwrap();
        println!("Parsed expression: {:?}", ast);

        let result = interp("scaled_value(4)", Some(Rc::new(ctx.clone())));
        match &result {
            Ok(val) => assert_eq!(*val, 10.0 + 4.0 * 2.5),
            Err(e) => {
                println!("Error evaluating scaled_value(4): {:?}", e);
                println!("Context variables at error: {:?}", ctx.variables);
                println!("Context constants at error: {:?}", ctx.constants);
                println!("Full context: variables={:?}, constants={:?}, expression_functions={:?}", ctx.variables, ctx.constants, ctx.function_registry.expression_functions.keys().collect::<Vec<_>>());
                panic!("Failed to evaluate scaled_value(4): {:?}", e);
            }
        }
    }

    #[test]
    fn test_recursive_expression_function() {
        let mut ctx = EvalContext::new();

        // Register a recursive function that calculates factorial
        let result1 = ctx.register_expression_function(
            "factorial",
            &["n"],
            "n <= 1 ? 1 : n * factorial(n - 1)",
        );
        assert!(
            result1.is_err(),
            "Should reject expressions with comparison operators and ternary syntax"
        );

        // Register a non-recursive version instead
        let result2 = ctx.register_expression_function(
            "factorial",
            &["n"],
            "n * (n - 1) * (n - 2) * (n - 3) * (n - 4) + (n <= 4 ? 0 : factorial(n - 5))",
        );
        assert!(
            result2.is_err(),
            "Should reject expressions with comparison operators and ternary syntax"
        );

        // Use a simpler approach with a limited factorial implementation
        let result3 = ctx.register_expression_function(
            "factorial5",
            &["n"],
            "n <= 1 ? 1 : n * (n - 1) * (n - 2) * (n - 3) * (n - 4) / 24 * 120",
        );
        assert!(
            result3.is_err(),
            "Should reject expressions with comparison operators and ternary syntax"
        );

        // Finally, use a non-recursive approach that works with our parser
        let result4 = ctx.register_expression_function(
            "factorial",
            &["n"],
            "n * (n - 1) * (n - 2) * (n - 3) * (n - 4) * (n <= 5 ? 1 : 120)",
        );
        assert!(
            result4.is_err(),
            "Should reject expressions with comparison operators"
        );

        // Register a simple non-recursive factorial implementation that works with our parser
        ctx.register_expression_function(
            "factorial",
            &["n"],
            "n * (n - 1) * (n - 2) * (n - 3) * (n - 4)",
        )
        .unwrap();

        // Test the factorial function for n=5
        let result = interp("factorial(5)", Some(Rc::new(ctx.clone()))).unwrap();
        assert_eq!(result, 120.0); // 5! = 120

        // Register an extended factorial for n=6
        ctx.register_expression_function(
            "factorial6",
            &["n"],
            "n * (n - 1) * (n - 2) * (n - 3) * (n - 4) * (n - 5)",
        )
        .unwrap();

        // Test for n=6
        let result2 = interp("factorial6(6)", Some(Rc::new(ctx.clone()))).unwrap();
        assert_eq!(result2, 720.0); // 6! = 720
    }

    #[test]
    fn test_expression_function_with_constants() {
        let mut ctx = EvalContext::new();

        // Register a function that calculates the area of a circle
        ctx.register_expression_function("circle_area", &["radius"], "pi * radius^2")
            .unwrap();

        // Register a function that calculates the volume of a sphere
        ctx.register_expression_function("sphere_volume", &["radius"], "(4/3) * pi * radius^3")
            .unwrap();

        // Test the circle area function
        let result = interp("circle_area(2)", Some(Rc::new(ctx.clone()))).unwrap();
        assert_approx_eq!(
            result, constants::PI * 4.0, constants::TEST_PRECISION
        );

        // Test the sphere volume function
        let result2 = interp("sphere_volume(3)", Some(Rc::new(ctx.clone()))).unwrap();
        let expected = (4.0 / 3.0) * constants::PI * 27.0;
        assert_approx_eq!(
            result2, expected, constants::TEST_PRECISION
        );
    }

    #[test]
    fn test_expression_function_error_handling() {
        let mut ctx = EvalContext::new();

        // Register a function that could cause division by zero
        ctx.register_expression_function("safe_divide", &["x", "y"], "x / y")
            .unwrap();

        // Test with valid input
        let result = interp("safe_divide(10, 2)", Some(Rc::new(ctx.clone()))).unwrap();
        assert_eq!(result, 5.0);

        // Test with division by zero
        let result2 = interp("safe_divide(10, 0)", Some(Rc::new(ctx.clone()))).unwrap();
        assert!(
            result2.is_infinite(),
            "Division by zero should return infinity"
        );

        // Register a function that handles the error case explicitly
        let result3 =
            ctx.register_expression_function("better_divide", &["x", "y"], "y == 0 ? 0 : x / y");
        assert!(
            result3.is_err(),
            "Should reject expressions with comparison operators and ternary syntax"
        );

        // Use a workaround with a very small denominator instead
        let result4 = ctx.register_expression_function(
            "better_divide",
            &["x", "y"],
            "x / (y + (y == 0) * 1e-10)",
        );
        assert!(
            result4.is_err(),
            "Should reject expressions with comparison operators"
        );

        // Register the max function as a native function since it's not available as an expression function
        ctx.register_native_function(
            "max",
            2,
            |args| {
                if args[0] > args[1] {
                    args[0]
                } else {
                    args[1]
                }
            },
        );

        // Use a simpler approach that works with our parser
        ctx.register_expression_function("better_divide", &["x", "y"], "x / max(y, 1e-10)")
            .unwrap();

        // Test with division by zero using the better function
        let result3 = interp("better_divide(10, 0)", Some(Rc::new(ctx.clone()))).unwrap();
        println!("better_divide(10, 0) = {}", result3); // Debug output
                                                        // When y is 0, we use max(0, 1e-10) which is 1e-10
                                                        // So the result is 10 / 1e-10 = 1e11
        #[cfg(feature = "f32")]
        assert_approx_eq!(
            result3, 1e11 as Real, 1e6 as Real // Cast literals to Real
        );

        #[cfg(not(feature = "f32"))]
        assert_approx_eq!(
            result3, 1e11 as Real, 1e6 as Real // Cast literals to Real
        );
    }

    // No longer need to import approx_eq as we're using assert_approx_eq! macro
}