aimx/expressions/

closure.rs

//! # Closure expression parsing and invocation.
//!
//! This module provides parsing and evaluation support for closure expressions in the AIMX language.
//! Closures are anonymous functions that can capture variables from their surrounding scope and are
//! denoted by the `=>` operator.
//!
//! ## Syntax
//!
//! AIMX supports two forms of closures:
//! - Single parameter: `param => expression`
//! - Double parameter: `(param1, param2) => expression`
//!
//! Closures are first-class values that can be stored in variables, passed as arguments to functions,
//! and returned from functions.
//!
//! ## Examples
//!
//! ```text
//! // Simple closure with single parameter
//! x => x * 2
//!
//! // Closure with two parameters
//! (x, y) => x + y
//!
//! // Using closures with array operations
//! nums.filter(x => x > 5)
//! nums.map((item, index) => item * index)
//! ```

use nom::{
    IResult,
    Parser,
    character::complete::{char, multispace0},
    bytes::tag,
};
use crate::{
    ContextLike,
    ExpressionLike,
    expressions::{Conditional, parse_conditional, parse_procedure, parse_identifier, Procedure},
    Primary,
    Value,
    Writer
};
use std::fmt;
use anyhow::Result;

/// Represents a closure (anonymous function) expression in the AIMX language.
///
/// Closures capture their environment and can be used as first-class values.
/// They are typically created using the `=>` operator and can have one or two parameters.
///
/// # Variants
///
/// - `One(String, Procedure)`: A closure with a single parameter
/// - `Two(String, String, Procedure)`: A closure with two parameters  
/// - `Primary(Box<Primary>)`: An optimized representation when the closure is a simple primary expression
///
/// # Examples
///
/// ```text
/// // Parses to Closure::One("x", procedure)
/// x => x + 1
///
/// // Parses to Closure::Two("x", "y", procedure)
/// (x, y) => x + y
/// ```
#[derive(Debug, Clone, PartialEq)]
pub enum Closure {
    One(String, Procedure),
    Two(String, String, Procedure),  
    /// Primary flattened AST optimization - used when the closure body is a simple primary expression
    Primary(Box<Primary>),
}

impl Closure {
    /// Invokes the closure with the current context.
    ///
    /// This method evaluates the closure body after setting up the parameter bindings
    /// in the context. For closures with parameters, the parameters are bound to
    /// special context keys (0 for single parameter, 0 and 1 for double parameters).
    ///
    /// # Arguments
    ///
    /// * `context` - The evaluation context
    ///
    /// # Returns
    ///
    /// Returns `Result<Value>` containing the result of evaluating the closure body.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use aimx::{expressions::Closure, context::{Context, ContextLike}, expressions::Procedure};
    /// # use aimx::Primary;
    /// # use anyhow::Result;
    /// # fn example() -> Result<()> {
    /// let mut context = Context::new();
    /// let closure = Closure::One("x".to_string(), Procedure::Primary(Box::new(Primary::Literal(aimx::Literal::Number(42.0)))));
    /// 
    /// // Parameter values are set using the context's key mechanism
    /// let param_name = "x".to_string();
    /// context.set_key(0, &param_name);
    /// let result = closure.invoke(&mut context)?;
    /// // result should be Value::Literal(Literal::Number(42.0))
    /// # Ok(())
    /// # }
    /// ```
    pub fn invoke(&self, context: &mut dyn ContextLike) -> Result<Value> {
        match self {
            Closure::Two(one, two, conditional) => {
                context.set_key(0, one);
                context.set_key(1, two);
                conditional.evaluate(context)
            }
            Closure::One(one, conditional) => {
                context.set_key(0, one);
                conditional.evaluate(context)
            }
            Closure::Primary(primary) => primary.evaluate(context),       
        } 
    }
}

/// Parse a single-parameter closure expression.
///
/// Handles the syntax: `identifier => procedure`
fn single_closure(input: &str) -> IResult<&str, Closure> {
    let (input, one) = parse_identifier(input)?;
    let (input, _) = multispace0.parse(input)?;
    let (input, _) = tag("=>").parse(input)?;
    let (input, _) = multispace0.parse(input)?;
    let (input, procedure) = parse_procedure(input)?;
    Ok((input, Closure::One(one, procedure)))
}

/// Parse a two-parameter closure expression.
///
/// Handles the syntax: `(identifier, identifier) => procedure`
fn double_closure(input: &str) -> IResult<&str, Closure> {
    let (input, _) = char('(').parse(input)?;
    let (input, one) = parse_identifier(input)?;
    let (input, _) = multispace0.parse(input)?;
    let (input, _) = char(',').parse(input)?;
    let (input, _) = multispace0.parse(input)?;
    let (input, two) = parse_identifier(input)?;
    let (input, _) = multispace0.parse(input)?;
    let (input, _) = char(')').parse(input)?;
    let (input, _) = multispace0.parse(input)?;
    let (input, _) = tag("=>").parse(input)?;
    let (input, _) = multispace0.parse(input)?;
    let (input, procedure) = parse_procedure(input)?;
    Ok((input, Closure::Two(one, two, procedure)))
}

/// Parses a closure expression from the input string.
///
/// This is the main entry point for parsing closures. It attempts to parse
/// both single-parameter and two-parameter closures, falling back to
/// `parse_closure_conditional` if neither form matches.
///
/// # Arguments
///
/// * `input` - The input string to parse
///
/// # Returns
///
/// Returns `IResult<&str, Closure>` containing the remaining unparsed input
/// and the parsed closure.
///
/// # Examples
///
/// ```rust
/// # use aimx::expressions::closure::parse_closure;
/// let result = parse_closure("x => x * 2");
/// assert!(result.is_ok());
/// 
/// let result = parse_closure("(a, b) => a + b");
/// assert!(result.is_ok());
/// 
/// let result = parse_closure("simple_expression");
/// assert!(result.is_ok()); // Falls back to conditional parsing
/// ```
pub fn parse_closure(input: &str) -> IResult<&str, Closure> {
    if input.starts_with('(') {
        if let Ok((input, closure)) = double_closure(input) {
            return Ok((input, closure));
        }
    } else if let Ok((input, closure)) = single_closure(input) {
        return Ok((input, closure));
    }
    parse_closure_conditional(input)
}

/// Parses a closure that falls back to conditional expression parsing.
///
/// This function is used when the input doesn't match the standard closure
/// syntax. It parses the input as a conditional expression and wraps it
/// in a `Closure::Primary` variant.
///
/// # Arguments
///
/// * `input` - The input string to parse
///
/// # Returns
///
/// Returns `IResult<&str, Closure>` containing the remaining unparsed input
/// and the parsed closure.
pub fn parse_closure_conditional(input: &str) -> IResult<&str, Closure> {
    let (input, conditional) = parse_conditional(input)?;
    let closure = match conditional {
        Conditional::Primary(primary) => Closure::Primary(primary),
        _ => Closure::Primary(Box::new(Primary::Conditional(conditional))),
    };
    Ok((input, closure))
}

impl ExpressionLike for Closure {
    /// Evaluates the closure expression.
    ///
    /// For primary closures (optimized simple expressions), this evaluates the expression directly.
    /// For parameterized closures, this returns a `Value::Closure` containing the closure itself,
    /// since closures with parameters need to be invoked with `invoke()` to evaluate them.
    ///
    /// # Arguments
    ///
    /// * `context` - The evaluation context
    ///
    /// # Returns
    ///
    /// Returns `Result<Value>` containing either the evaluated result (for primary closures)
    /// or the closure itself as a value (for parameterized closures).
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use aimx::{Context, ExpressionLike};
    /// # use aimx::expressions::{closure::Closure, primary::Primary};
    /// # use anyhow::Result;
    /// # fn example() -> Result<()> {
    /// let mut context = Context::new();
    /// let closure = Closure::Primary(Box::new(Primary::Literal(aimx::Literal::Number(42.0))));
    /// let result = closure.evaluate(&mut context)?;
    /// // result should be Value::Literal(Literal::Number(42.0))
    /// # Ok(())
    /// # }
    /// ```
    fn evaluate(&self, context: &mut dyn ContextLike) -> Result<Value> {
        match self {
            Closure::Primary(primary) => primary.evaluate(context),
            _ => Ok(Value::Closure(self.clone())),
        }
    }

    /// Writes the closure expression to a writer.
    ///
    /// This method formats the closure according to its variant and writes it
    /// to the provided writer. The output matches the AIMX syntax for closures.
    ///
    /// # Arguments
    ///
    /// * `writer` - The writer to write the formatted closure to
    fn write(&self, writer: &mut Writer) {
        match self {
            Closure::One(one, conditional) => {
                writer.write_str(one);
                writer.write_str(" => ");
                conditional.write(writer);
            },
            Closure::Two(one, two, conditional) => {
                writer.write_char('(');
                writer.write_str(one);
                writer.write_str(", ");
                writer.write_str(two);
                writer.write_str(") => ");
                conditional.write(writer);
            },
            Closure::Primary(primary) => primary.write(writer),
        }
    }
    
    /// Returns a sanitized string representation of the closure.
    ///
    /// Sanitized output is suitable for display to users and may include
    /// formatting that makes the expression easier to read.
    ///
    /// # Returns
    ///
    /// Returns a `String` containing the sanitized closure representation.
    fn to_sanitized(&self) -> String {
        let mut writer = Writer::sanitizer();
        self.write(&mut writer);
        writer.finish()
    }
    
    /// Returns a formula string representation of the closure.
    ///
    /// Formula output is suitable for use as input to the parser and matches
    /// the exact AIMX syntax.
    ///
    /// # Returns
    ///
    /// Returns a `String` containing the formula representation.
    fn to_formula(&self) -> String {
        let mut writer = Writer::formulizer();
        self.write(&mut writer);
        writer.finish()
    }
}

impl fmt::Display for Closure {
    /// Formats the closure for display.
    ///
    /// This implementation uses the `Writable` trait to format the closure
    /// as a string suitable for user display.
    ///
    /// # Arguments
    ///
    /// * `f` - The formatter to write to
    ///
    /// # Returns
    ///
    /// Returns `fmt::Result` indicating success or failure.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut writer = Writer::stringizer();
        self.write(&mut writer);
        write!(f, "{}", writer.finish())
    }
}