rocks_lang/
environment.rs

1use std::fmt::Debug;
2use std::cell::RefCell;
3use std::collections::HashMap;
4use std::rc::Rc;
5
6use crate::object::Object;
7use crate::token::Token;
8use crate::error::{RuntimeError, Error};
9
10/// Represents an environment in which variables are stored.
11/// The environment is a hash map of variable names to their values.
12/// Each environment has a reference to its enclosing environment.
13/// This is an optional reference to implement lexical scoping and closures.
14#[derive(Clone)]
15pub struct Environment {
16    /// Using an Rc and Refcell here allows us to have multiple mutable references
17    /// to the same environment.
18    pub enclosing: Option<Rc<RefCell<Environment>>>,
19    variables: HashMap<String, Object>,
20}
21
22impl Environment {
23    /// Creates a new environment with the given enclosing environment.
24    pub fn new(enclosing: Option<Rc<RefCell<Environment>>>) -> Self {
25        Environment {
26            enclosing,
27            variables: HashMap::new(),
28        }
29    }
30
31    /// Defines a new variable in the environment with the given name and value.
32    pub fn define(&mut self, name: &str, value: Object) {
33        self.variables.insert(name.to_string(), value);
34    }
35
36    /// Accesses the ancestor environment at the given distance.
37    fn ancestor(&self, distance: usize) -> Rc<RefCell<Environment>> {
38        let parent = self.enclosing.clone()
39            .unwrap_or_else(|| panic!("enclosing environment to exist at depth {}", 1));
40        let mut environment = Rc::clone(&parent);
41
42        for i in 1..distance {
43            let parent = environment.borrow().enclosing.clone()
44                .unwrap_or_else(|| panic!("enclosing environment to exist at depth {}", i));
45            environment = Rc::clone(&parent);
46        }
47
48        environment
49    }
50
51    /// Assigns the given value to the variable with the given name.
52    /// If the variable is not define in this environment but is defined in an enclosing environment,
53    /// it will try to recursively assign the value to the variable in the enclosing environment.
54    /// If the variable is not defined in this environment or any enclosing environment, it will
55    /// throw a runtime error.
56    pub fn assign(&mut self, name: &Token, value: Object) {
57        if self.variables.contains_key(&name.lexeme) {
58            self.variables.insert(name.lexeme.clone(), value);
59            return;
60        }
61
62        if let Some(enclosing) = &mut self.enclosing {
63            enclosing.borrow_mut().assign(name, value);
64            return;
65        }
66
67        RuntimeError {
68            token: name.clone(),
69            message: format!("Undefined variable '{}'", name.lexeme),
70        }.throw();
71    }
72
73    /// Works like [`Environment::assign`] but assigns the value to the variable in the ancestor
74    /// environment at the given distance.
75    pub fn assign_at(&mut self, distance: usize, name: &Token, value: Object) {
76        if distance == 0 {
77            return self.assign(name, value);
78        }
79
80        self.ancestor(distance).borrow_mut().assign(name, value);
81    }
82
83    /// Returns the value of the variable with the given name.
84    /// If the variable is not defined in this environment but is defined in an enclosing environment,
85    /// it will try to recursively get the value of the variable in the enclosing environment.
86    /// If the variable is not defined in this environment or any enclosing environment, it will
87    /// throw a runtime error.
88    pub fn get(&self, name: &Token) -> Result<Object, RuntimeError> {
89        if let Some(variable) = self.variables.get(&name.lexeme) {
90            return Ok(variable.clone());
91        }
92
93        if let Some(enclosing) = &self.enclosing {
94            return enclosing.borrow().get(name);
95        }
96
97        Err(RuntimeError {
98            token: name.clone(),
99            message: format!("Undefined variable '{}'", name.lexeme)
100        })
101    }
102
103    /// Works like [`Environment::get`] but gets the value of the variable in the ancestor
104    /// environment at the given distance.
105    pub fn get_at(&self, distance: usize, name: &Token) -> Result<Object, RuntimeError> {
106        if distance == 0 {
107            return self.get(name);
108        }
109
110        return self.ancestor(distance).borrow().get(name);
111    }
112}
113
114impl Default for Environment {
115    fn default() -> Self {
116        Self::new(None)
117    }
118}
119
120impl Debug for Environment {
121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        f.debug_struct("Environment")
123            .field("enclosing", &self.enclosing)
124            .field("variables", &self.variables.keys())
125            .finish()
126    }
127}