sunflower/src/expression.rs

use std::str::FromStr;

use crate::{env::Env, function::FunctionMap, statement::Statement, value::Value};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BinaryOperation {
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    LessThan,
    GreaterThan,
    LessThanOrEqual,
    GreaterThanOrEqual,
    Equal,
    NotEqual,
}

impl FromStr for BinaryOperation {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "+" => Ok(BinaryOperation::Add),
            "-" => Ok(BinaryOperation::Sub),
            "*" => Ok(BinaryOperation::Mul),
            "/" => Ok(BinaryOperation::Div),
            "%" => Ok(BinaryOperation::Mod),
            "<" => Ok(BinaryOperation::LessThan),
            ">" => Ok(BinaryOperation::GreaterThan),
            "<=" => Ok(BinaryOperation::LessThanOrEqual),
            ">=" => Ok(BinaryOperation::GreaterThanOrEqual),
            "==" => Ok(BinaryOperation::Equal),
            "!=" => Ok(BinaryOperation::NotEqual),
            _ => Err("Invalid binary operation"),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Expression {
    Constant(Value),
    Variable(String),
    Call(String, Vec<Expression>),
    CallMethod(Box<Expression>, String, Vec<Expression>),
    Binary(Box<Expression>, BinaryOperation, Box<Expression>),
    If(
        Vec<(Expression, Statement, Expression)>,
        (Box<Statement>, Box<Expression>),
    ),
}

impl Expression {
    pub fn eval(&self, e: &Env, f: &FunctionMap) -> Value {
        match self {
            Expression::Constant(v) => v.clone(),
            Expression::Variable(name) => e.get(name).unwrap_or_default(),
            Expression::Call(name, args) => {
                let args = args.iter().map(|ex| ex.eval(e, f)).collect::<Vec<_>>();
                let res = f.call(name, e, &args);
                if let Ok(res) = res {
                    res
                } else {
                    args[0].call_method(name, e, f, &args)
                }
            }
            Expression::CallMethod(target, name, args) => {
                let target = target.eval(e, f);
                let v = vec![target.clone()];
                let args = v
                    .iter()
                    .cloned()
                    .chain(args.iter().map(|ex| ex.eval(e, f)))
                    .collect::<Vec<_>>();

                target.call_method(name, e, f, &args)
            }
            Expression::Binary(lhs, op, rhs) => {
                let lhs = lhs.eval(e, f);
                let rhs = rhs.eval(e, f);
                match op {
                    BinaryOperation::Add => lhs + rhs,
                    BinaryOperation::Sub => lhs - rhs,
                    BinaryOperation::Mul => lhs * rhs,
                    BinaryOperation::Div => lhs / rhs,
                    BinaryOperation::Mod => lhs % rhs,
                    BinaryOperation::LessThan => (lhs < rhs).into(),
                    BinaryOperation::GreaterThan => (lhs > rhs).into(),
                    BinaryOperation::LessThanOrEqual => (lhs <= rhs).into(),
                    BinaryOperation::GreaterThanOrEqual => (lhs >= rhs).into(),
                    BinaryOperation::Equal => (lhs == rhs).into(),
                    BinaryOperation::NotEqual => (lhs != rhs).into(),
                }
            }
            Expression::If(ifs, (else_statement, else_expr)) => {
                for (expr, statement, ret) in ifs {
                    let e = Env::new_with_parent(e);
                    if expr.eval(&e, f).is_truthy() {
                        statement.eval(&e, f);
                        return ret.eval(&e, f);
                    }
                }

                let e = Env::new_with_parent(e);
                else_statement.eval(&e, f);
                else_expr.eval(&e, f)
            }
        }
    }
}

impl From<Value> for Expression {
    fn from(v: Value) -> Self {
        Expression::Constant(v)
    }
}

macro_rules! from_impl {
    ($($type:ty),+) => {
        $(
            impl From<$type> for Expression {
                fn from(v: $type) -> Self {
                    Expression::Constant(Value::from(v))
                }
            }
        )+
    };
}

from_impl!(i32, f32, String, bool, char);

impl<'s> From<&'s str> for Expression {
    fn from(s: &'s str) -> Self {
        Expression::Constant(Value::String(s.to_string()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::env::Env;
    use crate::value::Value;

    #[test]
    fn test_eval() {
        let f = FunctionMap::new();
        let e = Env::new();
        e.set("test", 10);
        assert_eq!(
            Expression::Variable("test".to_string()).eval(&e, &f),
            Value::Int(10)
        );
        assert_eq!(
            Expression::Constant(Value::Int(10)).eval(&e, &f),
            Value::Int(10)
        );
        assert_eq!(
            Expression::Binary(
                Box::new(Expression::Variable("test".to_string())),
                BinaryOperation::Add,
                Box::new(Expression::Constant(Value::Int(10)))
            )
            .eval(&e, &f),
            Value::Int(20)
        );
    }

    #[test]
    fn test_eval_unset() {
        let f = FunctionMap::new();

        let e = Env::new();
        assert_eq!(
            Expression::Variable("test".to_string()).eval(&e, &f),
            Value::Nothing
        );
        assert_eq!(
            Expression::Constant(Value::Error).eval(&e, &f),
            Value::Error
        );
        assert_eq!(
            Expression::Binary(
                Box::new(Expression::Variable("test".to_string())),
                BinaryOperation::Add,
                Box::new(Expression::Constant(Value::Nothing))
            )
            .eval(&e, &f),
            Value::Error
        );
    }

    #[test]
    fn test_eval_int() {
        let f = FunctionMap::new();

        let e = Env::new();
        e.set("test", 10);
        assert_eq!(
            Expression::Constant(Value::Int(10)).eval(&e, &f),
            Value::Int(10)
        );
        assert_eq!(
            Expression::Binary(
                Box::new(Expression::Constant(Value::Int(10))),
                BinaryOperation::Add,
                Box::new(Expression::Variable("test".to_string()))
            )
            .eval(&e, &f),
            Value::Int(20)
        );
    }

    #[test]
    fn eq_method_test() {
        let f = FunctionMap::new();
        let e = Env::new();
        e.set("test", 10);
        assert_eq!(
            Expression::CallMethod(
                Box::new(Expression::Variable("test".to_string())),
                "eq".to_string(),
                vec![Expression::Constant(Value::Int(10))]
            )
            .eval(&e, &f),
            Value::Bool(true)
        );
        assert_eq!(
            Expression::CallMethod(
                Box::new(Expression::Variable("test".to_string())),
                "eq".to_string(),
                vec![Expression::Constant(Value::Int(11))]
            )
            .eval(&e, &f),
            Value::Bool(false)
        );
    }
}