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

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

use thiserror::Error;

use crate::{token::Token, value::Value};

#[derive(Error, Debug)]
pub enum EnvironmentError {
    #[error("line {0}: undefined variable: {1}")]
    UndefinedVariable(usize, String),
}

/// Environment mapping variable names to their respective values. Can have an optional enclosing
/// environment. The outermost (global) environment has no enclosing one.
#[derive(Default, Debug, Clone)]
pub struct Environment {
    values: HashMap<String, Value>,
    enclosing: Option<Rc<RefCell<Environment>>>,
}

impl Environment {
    /// Initialize a new environment with an enclosing one.
    pub fn with_enclosing(enclosing: Rc<RefCell<Environment>>) -> Self {
        Self {
            values: HashMap::default(),
            enclosing: Some(enclosing),
        }
    }

    /// Define a new variable binding with a value.
    ///
    /// Overwrite any binding of the same name.
    pub fn define(&mut self, name: String, value: Value) {
        self.values.insert(name, value);
    }

    /// Assign a new value to an existing (defined) variable. Error if there is no such variable in
    /// this environment or any of the enclosing ones.
    pub fn assign(&mut self, token: &Token, value: Value) -> Result<(), EnvironmentError> {
        if self.values.contains_key(token.lexeme.as_str()) {
            self.values.insert(token.lexeme.clone(), value);
            Ok(())
        } else if let Some(enclosing) = &mut self.enclosing {
            enclosing.borrow_mut().assign(token, value)
        } else {
            Err(EnvironmentError::UndefinedVariable(
                token.line,
                token.lexeme.clone(),
            ))
        }
    }

    /// Get the value of an existing (defined) variable. Error if there is no such variable in
    /// this environment or any of the enclosing ones.
    pub fn get(&self, token: &Token) -> Result<Value, EnvironmentError> {
        if let Some(v) = self.values.get(token.lexeme.as_str()) {
            Ok(v.clone())
        } else if let Some(enclosing) = &self.enclosing {
            enclosing.borrow().get(token)
        } else {
            Err(EnvironmentError::UndefinedVariable(
                token.line,
                token.lexeme.clone(),
            ))
        }
    }
}