rlox/interpreter/src/interpreter.rs

use std::rc::Rc;

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

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> {
        match self {
            Literal::String(s) => Ok(Value::String(s.clone())),
            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 op = *op;

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

                if op == BinaryOp::Equal {
                    return Ok(Bool(left == right));
                }

                if op == BinaryOp::NotEqual {
                    return Ok(Bool(left != right));
                }

                match (left, right) {
                    (Number(left), Number(right)) => match op {
                        BinaryOp::Add => Ok(Number(left + right)),
                        BinaryOp::Subtract => Ok(Number(left - right)),
                        BinaryOp::Multiply => Ok(Number(left * right)),
                        BinaryOp::Divide => Ok(Number(left / right)),
                        BinaryOp::Less => Ok(Bool(left < right)),
                        BinaryOp::LessEqual => Ok(Bool(left <= right)),
                        BinaryOp::Greater => Ok(Bool(left > right)),
                        BinaryOp::GreaterEqual => Ok(Bool(left >= right)),
                        _ => unreachable!(),
                    },
                    (String(left), String(right)) => match op {
                        BinaryOp::Add => {
                            let s: Box<str> =
                                itertools::chain(left.as_ref().chars(), right.chars()).collect();

                            Ok(Value::string(s))
                        }
                        BinaryOp::Less => Ok(Bool(left < right)),
                        BinaryOp::LessEqual => Ok(Bool(left <= right)),
                        BinaryOp::Greater => Ok(Bool(left > right)),
                        BinaryOp::GreaterEqual => Ok(Bool(left >= right)),
                        BinaryOp::Equal | BinaryOp::NotEqual => unreachable!(),
                        _ => Err(RuntimeError::BinaryOpInvalidArguments {
                            left: String(left),
                            op,
                            right: String(right),
                        }),
                    },
                    (left, right) => {
                        Err(RuntimeError::BinaryOpInvalidArguments { left, op, right })
                    }
                }
            }
            Expr::Logical { left, op, right } => {
                let left = left.eval(env)?;

                // shortcircuit
                if *op == LogicalOp::Or && left.is_truthy()
                    || *op == LogicalOp::And && !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)?;

                for arg in args {
                    let arg = arg.eval(env)?;
                    env.push_arg(arg);
                }

                match callee {
                    Value::Function(fun) => call_fun(fun, env),
                    Value::ExternFunction(ext_fun) => call_extfun(ext_fun, env),
                    Value::Class(class) => call_class(class, 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.clone();

                        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.clone(), value);
                    Ok(Value::Nil)
                } else {
                    Err(RuntimeError::InvalidSetTarget)
                }
            }
            Expr::Function {
                name,
                param_names,
                closure_vars,
                body,
            } => Ok(Value::function(LoxFunction::new(
                name.clone(),
                env.collect_closure(closure_vars),
                param_names.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: FxHashMap<String, Value> = FxHashMap::default();

                // 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.clone(),
                    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 } => {
                use std::io::Write;

                let val = expr.eval(env)?;
                writeln!(env.runtime(), "{val}").unwrap();
            }
            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.clone(), 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)?;
            }
            Stmt::Break => return Err(RuntimeError::Break),
            Stmt::Return { expr } => {
                let value = expr.eval(env)?;
                return Err(RuntimeError::Return { value });
            }
        }

        Ok(Value::Nil)
    }
}

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

pub fn call_fun(fun: Rc<LoxFunction>, env: &mut Environment) -> EvalResult<Value> {
    let args = env.collect_args();

    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.clone(), 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
}

pub fn call_extfun(ext_fun: Rc<LoxExternFunction>, env: &mut Environment) -> EvalResult<Value> {
    let args = env.collect_args();

    if args.len() != ext_fun.arity() {
        return Err(RuntimeError::WrongArity {
            name: ext_fun.name().to_owned(),
            arity: ext_fun.arity(),
            given: args.len(),
        });
    }

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

// has to take class as an argument instead of as self to leave it behind its Rc
pub fn call_class(class: Rc<LoxClass>, env: &mut Environment) -> EvalResult<Value> {
    // let args = env.collect_args();

    if env.num_args() != class.arity() {
        return Err(RuntimeError::WrongArity {
            name: class.name().to_owned(),
            arity: class.arity(),
            given: env.num_args(),
        });
    }

    let object = LoxReference::new(class);

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

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

    Ok(Value::Object(object))
}