crafting-interpreters/rust/rox/src/interpreter.rs

use std::{cell::RefCell, collections::HashMap, rc::Rc};

use ordered_float::OrderedFloat;
use thiserror::Error;
use tracing::error;

use crate::{
    callable::CallingError,
    environment::{Environment, EnvironmentError},
    expression::Expression,
    function::Function,
    native_functions,
    statement::Statement,
    token::{Literal, Token, TokenType},
    value::Value,
};

#[derive(Error, Debug)]
pub enum InterpreterError {
    #[error("line {0}: MINUS unary expression expects a number on the right")]
    UnaryExpressionNotANumber(usize),
    #[error("line {0}: unknown unary operator: {1}")]
    UnaryOperatorUnknown(usize, String),
    #[error("line {0}: unknown binary operator: {1}")]
    BinaryOperatorUnknown(usize, String),
    #[error("line {0}: left or right is not a number.")]
    BinaryExpressionNeedsNumber(usize),
    #[error("line {0}: left or right is neither a number nor string.")]
    BinaryExpressionNeedsNumberOrString(usize),
    #[error("{0}")]
    UndefinedVariable(EnvironmentError),
    #[error("line {0}: {1} is not callable.")]
    NotACallable(usize, Value),
    #[error("line {0}: {1}.")]
    FailedToCall(usize, CallingError),
}

/// Interpreter for the Lox language.
#[derive(Debug)]
pub struct Interpreter {
    pub globals: Rc<RefCell<Environment>>,
    environment: Rc<RefCell<Environment>>,
    locals: HashMap<Expression, usize>,
}

/// Default configuration for the interpreter, with builtin native functions.
impl Default for Interpreter {
    fn default() -> Self {
        let env: Rc<RefCell<Environment>> = Default::default();

        {
            // add all native functions to global environment
            let mut env = env.borrow_mut();
            for (name, fun) in native_functions::all() {
                env.define(name, fun);
            }
        }

        Self {
            globals: env.clone(),
            environment: env.clone(),
            locals: HashMap::new(),
        }
    }
}

impl Interpreter {
    /// Try to evaluate an expression and return its result.
    pub fn run(&mut self, statements: Vec<Statement>) -> Result<(), InterpreterError> {
        for stmt in statements {
            match self.execute(&stmt) {
                Ok(_) => {}
                Err(e) => error!("{e}"),
            };
        }

        Ok(())
    }

    pub fn resolve(&mut self, expression: Expression, depth: usize) {
        self.locals.insert(expression, depth);
    }

    ///Execute a statement.
    fn execute(&mut self, statement: &Statement) -> Result<Option<Value>, InterpreterError> {
        match statement {
            Statement::Block(statements) => {
                let sub_env = Environment::with_enclosing(self.environment.clone());
                self.block(statements, sub_env)
            }
            Statement::Print(expression) => self.print_statement(expression),
            Statement::Expression(expression) => Ok(Some(self.evaluate(expression)?)),
            Statement::Var { name, initializer } => {
                self.var_statement(name, initializer.as_ref().as_ref())
            }
            Statement::If {
                condition,
                then_branch,
                else_branch,
            } => {
                let else_branch = else_branch.as_ref().map(|x| *x.clone());
                self.if_statement(condition, then_branch, else_branch.as_ref())
            }
            Statement::While { condition, body } => self.while_statement(condition, body),
            Statement::Function { name, params, body } => {
                self.function_statement(name, params, body)
            }
            Statement::Return { keyword: _, value } => self.return_statement(value),
        }
    }

    /// Execute all statements within a block, using a new environment (with the old one as the
    /// enclosing one). Immediately return the value if a Return Value is encountered after
    /// execution of a statement.
    pub fn block(
        &mut self,
        statements: &[Statement],
        environment: Environment,
    ) -> Result<Option<Value>, InterpreterError> {
        let prev_env = self.environment.clone();
        self.environment = Rc::new(RefCell::new(environment));

        for stmt in statements {
            let execution = self.execute(stmt);
            match execution {
                Ok(x) => {
                    if let Some(Value::Return(x)) = x {
                        self.environment = prev_env.clone();
                        return Ok(Some(*x));
                    }
                }
                Err(e) => error!("{e}"),
            }
        }

        self.environment = prev_env.clone();

        Ok(None)
    }

    /// Evaluate an expression and return its value.
    fn evaluate(&mut self, expression: &Expression) -> Result<Value, InterpreterError> {
        match expression {
            Expression::Literal { value } => self.literal(value.clone()),
            Expression::Grouping { expression } => self.grouping(expression),
            Expression::Unary {
                operator: op,
                right,
            } => self.unary(op, right),
            Expression::Binary {
                left,
                operator,
                right,
            } => self.binary(left, operator, right),
            Expression::Variable { name } => self.var_expression(name, expression),
            Expression::Assign { name, value } => self.assign(name, value),
            Expression::Logical {
                left,
                operator,
                right,
            } => self.logical_expression(left, operator, right),
            Expression::Call {
                callee,
                paren,
                args,
            } => self.call(callee, paren, args),
        }
    }

    /// Call a callable if it is one (meaning it starts with a LeftParen after an identifier),
    /// otherwise evaluate the expression.
    fn call(
        &mut self,
        callee: &Expression,
        paren: &Token,
        arg_expressions: &[Expression],
    ) -> Result<Value, InterpreterError> {
        let callee = self.evaluate(callee)?;
        let mut args = Vec::new();

        for arg in arg_expressions {
            args.push(self.evaluate(arg)?);
        }

        if let Value::Callable(f) = callee {
            f.call(self, args)
                .map_err(|e| InterpreterError::FailedToCall(paren.line, e))
        } else {
            Err(InterpreterError::NotACallable(paren.line, callee))
        }
    }

    /// Define a new function.
    fn function_statement(
        &mut self,
        name: &Token,
        params: &[Token],
        body: &[Statement],
    ) -> Result<Option<Value>, InterpreterError> {
        let fn_name = name.lexeme.clone();
        let fun = Function {
            name: name.clone(),
            params: params.to_vec(),
            body: body.to_vec(),
            closure: self.environment.clone(),
        };

        self.environment
            .borrow_mut()
            .define(fn_name, Value::Callable(Rc::new(fun)));

        Ok(None)
    }

    /// If the condition evaluates to truthy, execute the then branch, otherwise the else branch.
    fn if_statement(
        &mut self,
        condition: &Expression,
        then_branch: &Statement,
        else_branch: Option<&Statement>,
    ) -> Result<Option<Value>, InterpreterError> {
        if self.evaluate(condition)?.is_truthy() {
            self.execute(then_branch)
        } else if let Some(else_branch) = else_branch {
            self.execute(else_branch)
        } else {
            Ok(None)
        }
    }

    /// Evaluate an expression and print its value to stdout.
    fn print_statement(
        &mut self,
        expression: &Expression,
    ) -> Result<Option<Value>, InterpreterError> {
        let value = self.evaluate(expression)?;
        println!("{value}");

        Ok(None)
    }

    /// Return a value from a statement.
    fn return_statement(
        &mut self,
        value: &Option<Expression>,
    ) -> Result<Option<Value>, InterpreterError> {
        Ok(match value {
            Some(x) => Some(Value::Return(Box::new(self.evaluate(x)?))),
            None => Some(Value::Return(Box::new(Value::Nil))),
        })
    }

    /// Initialize a variable with an initializer expression or nil.
    fn var_statement(
        &mut self,
        name: &Token,
        initializer: Option<&Expression>,
    ) -> Result<Option<Value>, InterpreterError> {
        let value = if let Some(initializer) = initializer {
            self.evaluate(initializer)
        } else {
            Ok(Value::Nil)
        }?;

        self.environment
            .borrow_mut()
            .define(name.lexeme.clone(), value);

        Ok(None)
    }

    /// Execute the body as long as the condition evaluates to true.
    fn while_statement(
        &mut self,
        condition: &Expression,
        body: &Statement,
    ) -> Result<Option<Value>, InterpreterError> {
        while self.evaluate(condition)?.is_truthy() {
            self.execute(body)?;
        }

        Ok(None)
    }

    /// Assign the value of an expression to a variable.
    fn assign(&mut self, name: &Token, expression: &Expression) -> Result<Value, InterpreterError> {
        let value = self.evaluate(expression)?;

        self.environment
            .borrow_mut()
            .assign(name, value.clone())
            .map_err(InterpreterError::UndefinedVariable)?;

        if let Some(distance) = self.locals.get(expression) {
            self.environment
                .borrow_mut()
                .assign_at(*distance, name, value.clone())
                .map_err(InterpreterError::UndefinedVariable)?;
        } else {
            self.globals
                .borrow_mut()
                .assign(name, value.clone())
                .map_err(InterpreterError::UndefinedVariable)?;
        }

        Ok(value)
    }

    /// Convert the literal value into a Value.
    fn literal(&self, literal: Literal) -> Result<Value, InterpreterError> {
        Ok(literal.into())
    }

    /// Evaluate left and if the operator is Or and left is truthy or the operator is not Or and
    /// not leftis truthy short circuit and return left. Otherwise evaluate and return right.
    fn logical_expression(
        &mut self,
        left: &Expression,
        operator: &Token,
        right: &Expression,
    ) -> Result<Value, InterpreterError> {
        let left = self.evaluate(left)?;

        let truthy = if operator.token_type == TokenType::Or {
            left.is_truthy()
        } else {
            !left.is_truthy()
        };

        if truthy {
            Ok(left)
        } else {
            self.evaluate(right)
        }
    }

    /// Evaluate the inner expression.
    fn grouping(&mut self, inner: &Expression) -> Result<Value, InterpreterError> {
        self.evaluate(inner)
    }

    /// Evaluate the expression on the right and use its result when evaluating the unary operator.
    fn unary(&mut self, op: &Token, right: &Expression) -> Result<Value, InterpreterError> {
        let right = self.evaluate(right)?;

        match op.token_type {
            TokenType::Minus => {
                if let Value::Number(val) = right {
                    Ok(Value::Number(-val))
                } else {
                    Err(InterpreterError::UnaryExpressionNotANumber(op.line))
                }
            }
            TokenType::Bang => Ok(Value::Boolean(!right.is_truthy())),
            _ => Err(InterpreterError::UnaryOperatorUnknown(
                op.line,
                op.lexeme.clone(),
            )),
        }
    }

    /// Get the value of a variable.
    fn var_expression(
        &mut self,
        name: &Token,
        expression: &Expression,
    ) -> Result<Value, InterpreterError> {
        self.lookup_var(name, expression)
    }

    fn lookup_var(
        &mut self,
        name: &Token,
        expression: &Expression,
    ) -> Result<Value, InterpreterError> {
        if let Some(distance) = self.locals.get(expression) {
            self.environment
                .borrow()
                .get_at(*distance, name)
                .map_err(InterpreterError::UndefinedVariable)
        } else {
            self.globals
                .borrow()
                .get(name)
                .map_err(InterpreterError::UndefinedVariable)
        }
    }

    /// Calculate number operations.
    fn number_op(&self, left: f64, op: TokenType, right: f64) -> f64 {
        match op {
            TokenType::Minus => left - right,
            TokenType::Plus => left + right,
            TokenType::Slash => left / right,
            TokenType::Star => left * right,
            _ => unreachable!(),
        }
    }

    /// Calculate boolean operations.
    fn boolean_op(&self, left: f64, op: TokenType, right: f64) -> bool {
        match op {
            TokenType::Greater => left > right,
            TokenType::GreaterEqual => left >= right,
            TokenType::Less => left < right,
            TokenType::LessEqual => left <= right,
            _ => unreachable!(),
        }
    }

    /// Evaluate the left and right expressions (in that order) and then combine them with the
    /// specified operator.
    fn binary(
        &mut self,
        left: &Expression,
        op: &Token,
        right: &Expression,
    ) -> Result<Value, InterpreterError> {
        let left = self.evaluate(left)?;
        let right = self.evaluate(right)?;

        match op.token_type {
            TokenType::Minus | TokenType::Slash | TokenType::Star | TokenType::Plus => {
                if let (Value::Number(left), Value::Number(right)) = (left.clone(), right.clone()) {
                    Ok(Value::Number(OrderedFloat(self.number_op(
                        *left,
                        op.token_type,
                        *right,
                    ))))
                } else if let (Value::String(left), Value::String(right)) = (left, right) {
                    Ok(Value::String(format!("{}{}", left.clone(), right.clone())))
                } else {
                    Err(InterpreterError::BinaryExpressionNeedsNumberOrString(
                        op.line,
                    ))
                }
            }
            TokenType::Greater
            | TokenType::GreaterEqual
            | TokenType::Less
            | TokenType::LessEqual => {
                if let (Value::Number(left), Value::Number(right)) = (left, right) {
                    Ok(Value::Boolean(self.boolean_op(
                        *left,
                        op.token_type,
                        *right,
                    )))
                } else {
                    Err(InterpreterError::BinaryExpressionNeedsNumber(op.line))
                }
            }
            TokenType::BangEqual => Ok(Value::Boolean(left != right)),
            TokenType::EqualEqual => Ok(Value::Boolean(left == right)),
            _ => Err(InterpreterError::BinaryOperatorUnknown(
                op.line,
                op.lexeme.clone(),
            )),
        }
    }
}