[fix] New implementation for the 'lookup_variable' function

Author: rbonifacio

src/environment/environment.rs

@@ -28,9 +28,13 @@ impl<A: Clone> Scope<A> {
         return ();
     }
 
-    fn lookup(&self, var: &Name) -> Option<&A> {
+    fn lookup_var(&self, var: &Name) -> Option<&A> {
         self.variables.get(var)
     }
+
+    fn lookup_function(&self, name: &Name) -> Option<&Function> {
+        self.functions.get(name)
+    }
 }
 
 #[derive(Clone)]
@@ -51,20 +55,36 @@ impl<A: Clone> Environment<A> {
         match self.stack.front_mut() {
             None => self.globals.map_variable(var, value),
             Some(top) => top.map_variable(var, value),
-        };
-        return ();
+        }
     }
 
     pub fn map_function(&mut self, function: Function) -> () {
-        self.globals.map_function(function);
-        return ();
+        match self.stack.front_mut() {
+            None => self.globals.map_function(function),
+            Some(top) => top.map_function(function),
+        }
     }
 
     pub fn lookup(&self, var: &Name) -> Option<&A> {
-        match self.stack.front() {
-            None => self.globals.lookup(var),
-            Some(top) => top.lookup(var),
+        // First check local scopes in order
+        for scope in self.stack.iter() {
+            if let Some(value) = scope.lookup_var(var) {
+                return Some(value);
+            }
+        }
+        // Then check global scope
+        self.globals.lookup_var(var)
+    }
+
+    pub fn lookup_function(&self, name: &Name) -> Option<&Function> {
+        // First check local scopes in order
+        for scope in self.stack.iter() {
+            if let Some(func) = scope.lookup_function(name) {
+                return Some(func);
+            }
         }
+        // Then check global scope
+        self.globals.lookup_function(name)
     }
 
     pub fn scoped_function(&self) -> bool {
@@ -82,37 +102,71 @@ impl<A: Clone> Environment<A> {
 
 #[cfg(test)]
 mod tests {
-    use crate::environment::environment::{Environment, Scope};
+    use super::*;
 
     #[test]
-    fn eval_map_and_lookup_var() {
-        let mut s: Scope<i32> = Scope::new();
+    fn test_variable_scoping() {
+        let mut env: Environment<i32> = Environment::new();
+
+        // Test global scope
+        env.map_variable("x".to_string(), 32);
+        assert_eq!(Some(&32), env.lookup(&"x".to_string()));
 
-        s.map_variable("x".to_string(), 32);
-        s.map_variable("y".to_string(), 23);
+        // Test nested scopes
+        env.push();  // scope 1
+        env.map_variable("y".to_string(), 23);
+        env.map_variable("x".to_string(), 55);  // shadows global x
+
+        env.push();  // scope 2
+        env.map_variable("z".to_string(), 44);
 
-        assert_eq!(Some(32), s.lookup(&"x".to_string()).copied());
-        assert_eq!(Some(23), s.lookup(&"y".to_string()).copied());
+        // Variables from all scopes should be accessible
+        assert_eq!(Some(&55), env.lookup(&"x".to_string())); // from scope 1
+        assert_eq!(Some(&23), env.lookup(&"y".to_string())); // from scope 1
+        assert_eq!(Some(&44), env.lookup(&"z".to_string())); // from scope 2
+
+        // Pop scope 2
+        env.pop();
+        assert_eq!(Some(&55), env.lookup(&"x".to_string())); // still in scope 1
+        assert_eq!(Some(&23), env.lookup(&"y".to_string())); // still in scope 1
+        assert_eq!(None, env.lookup(&"z".to_string()));      // z is gone
+
+        // Pop scope 1
+        env.pop();
+        assert_eq!(Some(&32), env.lookup(&"x".to_string())); // back to global x
+        assert_eq!(None, env.lookup(&"y".to_string()));      // y is gone
     }
 
     #[test]
-    fn eval_environment() {
+    fn test_function_scoping() {
         let mut env: Environment<i32> = Environment::new();
+        
+        let global_func = Function {
+            name: "global".to_string(),
+            kind: None,
+            params: None,
+            body: None,
+        };
 
-        env.map_variable("x".to_string(), 32);
-
-        env.push();
+        let local_func = Function {
+            name: "local".to_string(),
+            kind: None,
+            params: None,
+            body: None,
+        };
 
-        env.map_variable("x".to_string(), 55);
-        env.map_variable("y".to_string(), 23);
+        // Test function scoping
+        env.map_function(global_func.clone());
+        assert!(env.lookup_function(&"global".to_string()).is_some());
 
-        assert_eq!(Some(55), env.lookup(&"x".to_string()).copied());
-        assert_eq!(Some(23), env.lookup(&"y".to_string()).copied());
-        assert_eq!(None, env.lookup(&"a".to_string()).copied());
+        env.push();
+        env.map_function(local_func.clone());
+        
+        assert!(env.lookup_function(&"global".to_string()).is_some()); // can see global
+        assert!(env.lookup_function(&"local".to_string()).is_some());  // can see local
 
         env.pop();
-
-        assert_eq!(Some(32), env.lookup(&"x".to_string()).copied());
-        assert_eq!(None, env.lookup(&"y".to_string()).copied());
+        assert!(env.lookup_function(&"global".to_string()).is_some()); // global still visible
+        assert!(env.lookup_function(&"local".to_string()).is_none());  // local gone
     }
 }