Feat: Implement Environment struct with variable management and integrate into Runtime

src/environment/mod.rs

@@ -0,0 +1,39 @@
+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

@@ -3,6 +3,7 @@ use std::panic;
 
 mod ast;
 mod compiler;
+mod environment;
 mod lexer;
 mod parser;
 mod vm;

src/vm/mod.rs

@@ -3,6 +3,7 @@ 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 {
@@ -34,7 +35,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 +56,38 @@ impl Runtime {
     }
 
     fn load_variable(&mut self, var: &str) {
-        if let Some(value) = self
-            .frames
-            .iter()
-            .rev()
-            .find_map(|frame| frame.variables.get(var))
-        {
-            self.stack.push((*value).clone());
+        let current_frame = self.frames.last().unwrap();
+
+        if let Some(value) = current_frame.env.borrow().get(var) {
+            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() {
+            let success = current_frame.env.borrow_mut().assign(&var, val);
+            if !success {
+                panic!("Variable '{}' not defined", var);
             }
         } else {
-            panic!("{} is not defined", var);
+            panic!("No defined value to store in {}", var);
         }
     }
 
@@ -214,8 +213,9 @@ 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 +315,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 +324,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
             }

tests/integration_tests.rs

@@ -1359,6 +1359,193 @@ print obj.inner["val"];
     assert_eq!(out, "5");
 }
 
+// ========== Closures / Lexical Scoping ==========
+
+#[test]
+fn test_function_mutates_outer_variable() {
+    let code = r#"
+let x = 10;
+function f() {
+    x = 20;
+}
+f();
+print x;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "20");
+}
+
+#[test]
+fn test_function_local_declaration_does_not_leak_to_outer() {
+    let code = r#"
+let x = 1;
+function f() {
+    let y = 99;
+    x = 2;
+}
+f();
+print x;
+"#;
+    let stderr = run_rts_should_fail(
+        r#"
+let x = 1;
+function f() {
+    let y = 99;
+}
+f();
+print y;
+"#,
+    );
+    assert!(stderr.contains("not defined"));
+
+    let out = run_rts(code);
+    assert_eq!(out, "2");
+}
+
+#[test]
+fn test_parameter_shadows_outer_variable() {
+    let code = r#"
+let x = 100;
+function f(x) {
+    x = x + 1;
+    return x;
+}
+print f(5);
+print x;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "6\n100");
+}
+
+#[test]
+fn test_local_let_shadows_outer_then_assign_targets_local() {
+    let code = r#"
+let x = 100;
+function f() {
+    let x = 1;
+    x = x + 1;
+    return x;
+}
+print f();
+print x;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "2\n100");
+}
+
+#[test]
+fn test_multiple_calls_accumulate_in_outer_variable() {
+    let code = r#"
+let count = 0;
+function inc() {
+    count = count + 1;
+}
+inc();
+inc();
+inc();
+print count;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "3");
+}
+
+#[test]
+fn test_two_functions_share_outer_state() {
+    let code = r#"
+let n = 5;
+function double_it() {
+    n = n * 2;
+}
+function add_one() {
+    n = n + 1;
+}
+double_it();
+add_one();
+print n;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "11");
+}
+
+#[test]
+fn test_function_mutates_outer_in_loop() {
+    let code = r#"
+let total = 0;
+function add(v) {
+    total = total + v;
+}
+for (let i = 1; i < 5; i++) {
+    add(i);
+}
+print total;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "10");
+}
+
+#[test]
+fn test_assign_to_undeclared_inside_function_fails() {
+    let stderr = run_rts_should_fail(
+        r#"
+function f() {
+    y = 5;
+}
+f();
+"#,
+    );
+    assert!(stderr.contains("not defined"));
+}
+
+#[test]
+fn test_nested_function_call_sees_global_through_parent_chain() {
+    let code = r#"
+let g = 7;
+function inner() {
+    return g;
+}
+function outer() {
+    return inner();
+}
+print outer();
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "7");
+}
+
+#[test]
+fn test_recursive_function_each_call_has_own_local_scope() {
+    let code = r#"
+function count_down(n) {
+    if (n == 0) {
+        return 0;
+    }
+    print n;
+    return count_down(n - 1);
+}
+count_down(3);
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "3\n2\n1");
+}
+
+#[test]
+fn test_outer_mutation_visible_after_recursive_chain() {
+    let code = r#"
+let calls = 0;
+function rec(n) {
+    calls = calls + 1;
+    if (n == 0) {
+        return 0;
+    }
+    return rec(n - 1);
+}
+rec(4);
+print calls;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "5");
+}
+
 #[test]
 fn test_greater_than_equals() {
     let code = r#"