mathexpr 0.1.1

//! Bytecode compiler for mathexpr.
//!
//! This module compiles AST expressions into bytecode for efficient evaluation.

#[cfg(not(feature = "std"))]
use alloc::{string::String, vec::Vec};

use hashbrown::HashMap;

use crate::ast::{BinOp, BuiltinFn, Expr};
use crate::error::CompileError;

/// A bytecode instruction for the expression VM.
#[derive(Debug, Clone)]
pub enum Instruction {
    /// Load a constant value onto the stack.
    LoadConst(f64),
    /// Load the current/input value onto the stack.
    LoadCurrent,
    /// Load a variable by index onto the stack.
    LoadVar(usize),
    /// Negate the top of the stack.
    Negate,
    /// Apply a binary operation to the top two stack values.
    BinOp(BinOp),
    /// Call a builtin function.
    Call(BuiltinFn),
}

/// A compiled expression ready for evaluation.
///
/// This contains the bytecode instructions and metadata needed to evaluate
/// the expression with different variable values.
#[derive(Debug, Clone)]
pub struct CompiledExpr {
    pub(crate) instructions: Vec<Instruction>,
    pub(crate) variable_indices: HashMap<String, usize>,
    pub(crate) uses_current: bool,
}

impl CompiledExpr {
    /// Compile an AST expression into bytecode.
    ///
    /// # Arguments
    ///
    /// * `expr` - The parsed AST expression
    /// * `var_names` - The variable names in the order they will be provided during evaluation
    ///
    /// # Returns
    ///
    /// * `Ok(CompiledExpr)` - The compiled bytecode
    /// * `Err(CompileError)` - If compilation fails (e.g., undefined variable, unknown function)
    ///
    /// # Example
    ///
    /// ```
    /// use mathexpr::{parse, CompiledExpr};
    ///
    /// let ast = parse("x + y").unwrap();
    /// let compiled = CompiledExpr::compile(&ast, &["x", "y"]).unwrap();
    /// ```
    pub fn compile(expr: &Expr, var_names: &[&str]) -> Result<Self, CompileError> {
        let variable_indices: HashMap<String, usize> = var_names
            .iter()
            .enumerate()
            .map(|(i, &name)| (String::from(name), i))
            .collect();

        let mut instructions = Vec::new();
        let mut uses_current = false;
        Self::compile_expr(
            expr,
            &variable_indices,
            &mut instructions,
            &mut uses_current,
        )?;

        Ok(CompiledExpr {
            instructions,
            variable_indices,
            uses_current,
        })
    }

    fn compile_expr(
        expr: &Expr,
        var_indices: &HashMap<String, usize>,
        instructions: &mut Vec<Instruction>,
        uses_current: &mut bool,
    ) -> Result<(), CompileError> {
        match expr {
            Expr::Number(n) => {
                instructions.push(Instruction::LoadConst(*n));
            }
            Expr::CurrentValue => {
                *uses_current = true;
                instructions.push(Instruction::LoadCurrent);
            }
            Expr::Variable(name) => {
                if let Some(&idx) = var_indices.get(name) {
                    instructions.push(Instruction::LoadVar(idx));
                } else {
                    return Err(CompileError::UndefinedVariable(name.clone()));
                }
            }
            Expr::UnaryMinus(e) => {
                Self::compile_expr(e, var_indices, instructions, uses_current)?;
                instructions.push(Instruction::Negate);
            }
            Expr::BinaryOp { op, left, right } => {
                Self::compile_expr(left, var_indices, instructions, uses_current)?;
                Self::compile_expr(right, var_indices, instructions, uses_current)?;
                instructions.push(Instruction::BinOp(*op));
            }
            Expr::FunctionCall { name, args } => {
                // Look up function at compile time
                if let Some((func, _arity)) = BuiltinFn::from_name(name) {
                    func.check_arity(args.len(), name)?;

                    // Compile all arguments (they'll be on the stack)
                    for arg in args {
                        Self::compile_expr(arg, var_indices, instructions, uses_current)?;
                    }
                    instructions.push(Instruction::Call(func));
                } else {
                    return Err(CompileError::UnknownFunction(name.clone()));
                }
            }
        }
        Ok(())
    }

    /// Get the number of variables expected during evaluation.
    pub fn num_variables(&self) -> usize {
        self.variable_indices.len()
    }

    /// Check if this expression uses the current/input value (`_`).
    pub fn uses_current_value(&self) -> bool {
        self.uses_current
    }

    /// Get the expected variable names in order.
    #[cfg(feature = "std")]
    pub fn variable_names(&self) -> Vec<&str> {
        let mut names: Vec<_> = self.variable_indices.iter().collect();
        names.sort_by_key(|(_, &idx)| idx);
        names.into_iter().map(|(name, _)| name.as_str()).collect()
    }
}