rlox/interpreter/src/interpret.rs

use std::collections::HashMap;
use std::rc::Rc;

use rlox2_frontend::parser::{BinaryOp, Expr, Literal, LogicalOp, Stmt, UnaryOp};

use crate::error::RuntimeError;
use crate::function::LoxExternFunction;
use crate::{LoxClass, LoxReference};

use super::environment::Environment;
use super::{LoxFunction, Runtime, Value};

pub type EvalResult<T> = Result<T, RuntimeError>;

/*====================================================================================================================*/

pub fn execute(statement: Stmt, runtime: &mut Runtime) -> Result<(), RuntimeError> {
    let mut env = Environment::new(runtime);

    statement.eval(&mut env)?;

    Ok(())
}

/*====================================================================================================================*/

trait Eval {
    fn eval(&self, env: &mut Environment) -> EvalResult<Value>;
}

impl Eval for Literal {
    fn eval(&self, env: &mut Environment) -> EvalResult<Value> {
        let _ = env;
        match self {
            Literal::String(s) => Ok(Value::String(Rc::clone(s))),
            Literal::Number(num) => Ok(Value::Number(*num)),
            Literal::Bool(b) => Ok(Value::Bool(*b)),
            Literal::Nil => Ok(Value::Nil),
        }
    }
}

impl Eval for Expr {
    fn eval(&self, env: &mut Environment) -> EvalResult<Value> {
        match self {
            Expr::Literal { literal } => literal.eval(env),
            Expr::Unary { op, expr } => {
                let arg = expr.eval(env)?;

                match (*op, arg) {
                    (UnaryOp::Negate, Value::Number(num)) => Ok(Value::Number(-num)),
                    (UnaryOp::Not, Value::Bool(b)) => Ok(Value::Bool(!b)),
                    (UnaryOp::Not, primitive) => Ok(Value::Bool(!primitive.is_truthy())),
                    (op, arg) => Err(RuntimeError::UnaryOpInvalidArgument { op, arg }),
                }
            }
            Expr::Binary { left, op, right } => {
                use Value::{Bool, Number, String};

                let left = left.eval(env)?;
                let right = right.eval(env)?;

                match (left, *op, right) {
                    (Number(left), BinaryOp::Add, Number(right)) => Ok(Number(left + right)),
                    (Number(left), BinaryOp::Subtract, Number(right)) => Ok(Number(left - right)),
                    (Number(left), BinaryOp::Multiply, Number(right)) => Ok(Number(left * right)),
                    (Number(left), BinaryOp::Divide, Number(right)) => {
                        if right == 0.0 {
                            return Err(RuntimeError::DivisionByZero);
                        }
                        Ok(Number(left / right))
                    }

                    (String(left), BinaryOp::Add, String(right)) => {
                        let mut s = std::string::String::with_capacity(left.capacity() + right.capacity());
                        s += &left;
                        s += &right;
                        Ok(String(Rc::new(s)))
                    }

                    (left, BinaryOp::Equal, right) => Ok(Bool(left == right)),
                    (left, BinaryOp::NotEqual, right) => Ok(Bool(left != right)),

                    (Number(left), BinaryOp::Less, Number(right)) => Ok(Bool(left < right)),
                    (Number(left), BinaryOp::LessEqual, Number(right)) => Ok(Bool(left <= right)),
                    (Number(left), BinaryOp::Greater, Number(right)) => Ok(Bool(left > right)),
                    (Number(left), BinaryOp::GreaterEqual, Number(right)) => Ok(Bool(left >= right)),

                    (String(left), BinaryOp::Less, String(right)) => Ok(Bool(left < right)),
                    (String(left), BinaryOp::LessEqual, String(right)) => Ok(Bool(left <= right)),
                    (String(left), BinaryOp::Greater, String(right)) => Ok(Bool(left > right)),
                    (String(left), BinaryOp::GreaterEqual, String(right)) => Ok(Bool(left >= right)),

                    (left, op, right) => Err(RuntimeError::BinaryOpInvalidArguments { left, op, right }),
                }
            }
            Expr::Logical { left, op, right } => {
                let left = left.eval(env)?;
                match op {
                    LogicalOp::Or => {
                        if left.is_truthy() {
                            return Ok(left);
                        }
                    }
                    LogicalOp::And => {
                        if !left.is_truthy() {
                            return Ok(left);
                        }
                    }
                }
                let right = right.eval(env)?;
                Ok(right)
            }
            Expr::Grouping { expr } => expr.eval(env),
            Expr::Variable { name } => panic!("Unresolved variable {name}"),
            Expr::LocalVariable { name, level } => env.get_local(name, *level),
            Expr::GlobalVariable { name } => env.get_global(name),
            Expr::Assignment { target, value } => {
                let value = value.eval(env)?;

                match target.as_ref() {
                    Expr::LocalVariable { name, level } => env.assign(name, value.clone(), *level)?,
                    Expr::GlobalVariable { name } => env.assign_global(name, value.clone())?,
                    _ => panic!("Invalid assigment target {target}"),
                }
                Ok(value)
            }
            Expr::Call { callee, args } => {
                let callee = callee.eval(env)?;
                let args = args
                    .iter()
                    .map(|arg| arg.eval(env))
                    .collect::<EvalResult<Vec<Value>>>()?;

                match callee {
                    Value::Function(fun) => LoxFunction::call(fun, args, env),
                    Value::ExternFunction(ext_fun) => ext_fun.call(args, env),
                    Value::Class(class) => LoxClass::call(class, args, env),
                    _ => Err(RuntimeError::NotCallable { callee }),
                }
            }
            Expr::Get { target, name } => {
                let target = target.eval(env)?;

                if let Value::Object(object) = target {
                    object.get(name).ok_or_else(|| {
                        let class = object.class();
                        let name = name.to_owned();
                        RuntimeError::UndefinedAttribute { class, name }
                    })
                } else {
                    Err(RuntimeError::InvalidGetTarget)
                }
            }
            Expr::Set { target, name, value } => {
                let target = target.eval(env)?;

                if let Value::Object(mut object) = target {
                    let value = value.eval(env)?;
                    object.set(name, value);
                    Ok(Value::Nil)
                } else {
                    Err(RuntimeError::InvalidSetTarget)
                }
            }
            Expr::Function {
                name,
                param_names,
                closure_vars,
                body,
            } => Ok(Value::function(LoxFunction::new(
                name.as_ref(),
                env.collect_closure(closure_vars),
                param_names.as_ref().clone(),
                body.as_ref().clone(),
            ))),
            Expr::Class {
                superclass,
                name,
                methods: method_exprs,
            } => {
                let superclass = match superclass.as_ref().map(|expr| expr.eval(env)) {
                    Some(Ok(Value::Class(superclass))) => Some(superclass),
                    Some(Ok(value)) => return Err(RuntimeError::InvalidSuperclass { value }),
                    Some(Err(err)) => return Err(err),
                    None => None,
                };

                let mut methods: HashMap<String, Value> = HashMap::new();

                // this is the scope "this" will get injected in
                env.enter_scope();

                for method_expr in method_exprs.iter() {
                    let method = method_expr.eval(env)?;
                    if let Value::Function(ref fun) = method {
                        let name = fun.name().to_owned();
                        methods.insert(name, method);
                    }
                }

                env.exit_scope();

                Ok(Value::class(LoxClass::new(name, methods, superclass)))
            }
            Expr::This => panic!("Unresolved this"),
            Expr::Super {
                super_var,
                this_var,
                method,
            } => match (super_var.eval(env)?, this_var.eval(env)?) {
                (Value::Class(superclass), Value::Object(this)) => {
                    if let Some(method) = superclass.get_method(method, this) {
                        Ok(method)
                    } else {
                        Err(RuntimeError::InvalidGetTarget)
                    }
                }
                (super_val, this_val) => panic!("super evaluated to {super_val} and this evaluated to {this_val}"),
            },
        }
    }
}

impl Eval for Stmt {
    fn eval(&self, env: &mut Environment) -> EvalResult<Value> {
        match self {
            Stmt::Print { expr } => {
                match expr.eval(env)? {
                    // special case: when printing a string, drop the surrounding ""
                    Value::String(s) => println!("{s}"),
                    val => println!("{val}"),
                }
            }
            Stmt::IfStmt {
                condition,
                then_branch,
                else_branch,
            } => {
                let condition = condition.eval(env)?;
                if condition.is_truthy() {
                    then_branch.eval(env)?;
                } else if let Some(else_branch) = else_branch {
                    else_branch.eval(env)?;
                }
            }
            Stmt::While { condition, body } => {
                while condition.eval(env)?.is_truthy() {
                    match body.eval(env) {
                        Ok(_) => {}
                        Err(RuntimeError::Break) => break,
                        Err(err) => return Err(err),
                    }
                }
            }
            Stmt::VarDecl { name, initializer } => {
                let initializer = initializer.eval(env)?;
                env.define(name, initializer);
            }
            Stmt::Block { statements } => {
                env.enter_scope();
                for statement in statements {
                    if let Err(err) = statement.eval(env) {
                        env.exit_scope();
                        return Err(err);
                    }
                }
                env.exit_scope();
            }
            Stmt::ExprStmt { expr } => {
                // expr.eval(env)?;
                // Ok(Value::Nil)
                expr.eval(env)?;
            }
            Stmt::Break => return Err(RuntimeError::Break),
            Stmt::Return { expr } => {
                let value = expr.eval(env)?;
                return Err(RuntimeError::Return { value });
            }
        }

        Ok(Value::Nil)
    }
}

/*====================================================================================================================*/

impl LoxFunction {
    pub fn call(fun: Rc<LoxFunction>, args: Vec<Value>, env: &mut Environment) -> EvalResult<Value> {
        if args.len() != fun.arity() {
            return Err(RuntimeError::WrongArity {
                name: fun.name().to_owned(),
                arity: fun.arity(),
                given: args.len(),
            });
        }

        env.enter_scope();

        env.define(fun.name(), Value::Function(fun.clone()));

        env.insert_closure(fun.closure().clone());

        for (name, value) in std::iter::zip(fun.param_names(), args) {
            env.define(name, value);
        }

        let ret_val = match fun.body().eval(env) {
            Ok(_) => Ok(Value::Nil),
            Err(RuntimeError::Return { value }) => Ok(value),
            Err(err) => Err(err),
        };

        env.exit_scope();

        ret_val
    }
}

impl LoxExternFunction {
    pub fn call(&self, args: Vec<Value>, env: &mut Environment) -> EvalResult<Value> {
        if args.len() != self.arity() {
            return Err(RuntimeError::WrongArity {
                name: self.name().to_owned(),
                arity: self.arity(),
                given: args.len(),
            });
        }

        (self.closure())(args, env)
    }
}

impl LoxClass {
    // has to take class as an argument instead of as self to leave it behind its Rc
    pub fn call(class: Rc<LoxClass>, args: Vec<Value>, env: &mut Environment) -> EvalResult<Value> {
        if args.len() != class.arity() {
            return Err(RuntimeError::WrongArity {
                name: class.name().to_owned(),
                arity: class.arity(),
                given: args.len(),
            });
        }

        let object = LoxReference::new(class);

        // object.init(args, env)?;

        if let Some(Value::Function(method)) = object.get("init") {
            LoxFunction::call(method, args, env)?;
        }

        Ok(Value::Object(object))
    }
}