Feat: Implement Environment struct for variable management and update runtime to use it

Author: shubhiscoding

example/test.rts

@@ -1,29 +1,9 @@
-let x = (2 + 3) * 4;
-print "x = " + x;
+let x = 10;
 
-function max(a, b) {
-    if (a > b) {
-        return a;
-    } else {
-        return b;
-    }
+function f() {
+    x = 20;
 }
 
-print "max is " + max(x, 15);
+f();
 
-function factorial(n) {
-    if (n < 2) { return 1; }
-    return n * factorial(n - 1);
-}
-print "5! = " + factorial(5);
-
-for (let i = 0; i < 5; i++) {
-    if (i == 2) { continue; }
-    if (i == 4) { break; }
-    print i;
-}
-
-let name = "Rusty";
-if (name && x > 10) {
-    print name + " works!";
-}
+print x;

src/environment/mod.rs

@@ -0,0 +1,45 @@
+use std::{cell::RefCell, collections::HashMap, rc::Rc};
+
+use crate::ast::Value;
+
+#[derive(Debug, Clone)]
+pub struct Environment {
+    pub variables: HashMap<String, Value>,
+    pub parent: Option<Rc<RefCell<Environment>>>,
+}
+
+impl Environment {
+    pub fn new(
+        parent: Option<Rc<RefCell<Environment>>>
+    ) -> Self {
+        Environment {
+            variables: HashMap::new(),
+            parent,
+        }
+    }
+
+    pub fn get(&self, name: &str) -> Option<Value> {
+        if let Some(value) = self.variables.get(name) {
+            Some(value.clone())
+        } else if let Some(parent) = self.parent.as_ref() {
+            parent.borrow().get(name)
+        } else {
+            None
+        }
+    }
+
+    pub fn assign(
+        &mut self,
+        name: &str,
+        value: Value
+    ) -> bool {
+        if self.variables.contains_key(name) {
+            self.variables.insert(name.to_string(), value);
+            true
+        } else if let Some(parent) = self.parent.as_mut() {
+            parent.borrow_mut().assign(name, value)
+        }else{
+            false
+        }
+    }
+}

src/main.rs

@@ -6,6 +6,7 @@ mod compiler;
 mod lexer;
 mod parser;
 mod vm;
+mod environment;
 use compiler::compile_program;
 use lexer::tokenize;
 use parser::Parser;

src/vm/mod.rs

@@ -3,8 +3,10 @@ use std::{cell::RefCell, collections::HashMap, fmt, rc::Rc};
 use crate::{
     ast::{BinaryOperation, Value},
     compiler::{FunctionBytecode, Instruction, Program},
+    environment::{Environment},
 };
 
+
 impl fmt::Display for Value {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
@@ -34,7 +36,7 @@ impl fmt::Display for Value {
 
 #[derive(Debug, Clone)]
 pub struct Frame {
-    variables: HashMap<String, Value>,
+    env: Rc<RefCell<Environment>>,
     instructions: Vec<Instruction>,
     ip: usize,
 }
@@ -55,40 +57,35 @@ impl Runtime {
     }
 
     fn load_variable(&mut self, var: &str) {
-        if let Some(value) = self
-            .frames
-            .iter()
-            .rev()
-            .find_map(|frame| frame.variables.get(var))
+        let current_frame =
+            self.frames.last().unwrap();
+
+        if let Some(value) =
+            current_frame
+                .env
+                .borrow()
+                .get(var)
         {
-            self.stack.push((*value).clone());
+            self.stack.push(value.clone());
         } else {
             panic!("{} is not defined", var);
         }
     }
 
     fn store_variable(&mut self, var: String) {
         if let Some(value) = self.stack.pop() {
-            self.frames.last_mut().unwrap().variables.insert(var, value);
+            self.frames.last_mut().unwrap().env.borrow_mut().variables.insert(var, value);
         } else {
             panic!("No defined value to store in {}", var);
         }
     }
 
     fn assign_variable(&mut self, var: String) {
-        if let Some(value) = self
-            .frames
-            .iter_mut()
-            .rev()
-            .find_map(|frame| frame.variables.get_mut(&var))
-        {
-            if let Some(val) = self.stack.pop() {
-                *value = val;
-            } else {
-                panic!("No defined value to store in {}", var);
-            }
+        let current_frame = self.frames.last().unwrap();
+        if let Some(val) = self.stack.pop() {
+            current_frame.env.borrow_mut().assign(&var, val);
         } else {
-            panic!("{} is not defined", var);
+            panic!("No defined value to store in {}", var);
         }
     }
 
@@ -214,8 +211,12 @@ impl Runtime {
 
 pub fn execute(program: Program, runtime: &mut Runtime) {
     runtime.functions = program.functions;
+    let global_env =
+    Rc::new(RefCell::new(
+        Environment::new(None)
+    ));
     runtime.frames.push(Frame {
-        variables: HashMap::new(),
+        env: global_env,
         instructions: program.main,
         ip: 0,
     });
@@ -315,7 +316,7 @@ pub fn execute(program: Program, runtime: &mut Runtime) {
             }
             Instruction::CallFunction(name, arg_count) => {
                 let func = runtime.functions.get(&name).expect("undefined function");
-                let mut locals = HashMap::new();
+                // let mut locals = HashMap::new();
                 // Pop args in reverse (last arg was pushed last)
                 if arg_count != func.params.len() {
                     panic!(
@@ -324,16 +325,21 @@ pub fn execute(program: Program, runtime: &mut Runtime) {
                         arg_count
                     );
                 }
+                let parnt_env = runtime.frames.last().unwrap().env.clone();
+                let local_env =
+                    Rc::new(RefCell::new(
+                        Environment::new(Some(parnt_env))
+                    ));
                 for param in func.params.iter().rev() {
-                    locals.insert(param.clone(), runtime.stack.pop().unwrap());
+                    local_env.borrow_mut().variables.insert(param.clone(), runtime.stack.pop().unwrap());
                 }
                 // Save caller's position (move past the CallFunction instruction)
                 runtime.frames.last_mut().unwrap().ip += 1;
                 // Push new frame
                 runtime.frames.push(Frame {
                     instructions: func.instructions.clone(),
                     ip: 0,
-                    variables: locals,
+                    env: local_env
                 });
                 continue; // don't increment ip again
             }