aster-core/src/vm/compiler.rs

@@ -97,11 +97,16 @@ struct LoopContext {
 pub struct Compiler {
     function: FunctionProto,
     locals: Vec<Local>,
-    upvalues: Vec<Upvalue>,
-    /// Direct parent's locals (for upvalue resolution)
+    /// Upvalues for this function (shared with child for transitive resolution)
+    upvalues: Rc<RefCell<Vec<Upvalue>>>,
+    /// ALL enclosing locals from entire chain (merged, direct parent first)
     enclosing_locals: Option<Vec<Local>>,
-    /// Direct parent's upvalues (for transitive upvalue resolution)
-    enclosing_upvalues: Option<Vec<Upvalue>>,
+    /// Direct parent's own locals count (to distinguish parent locals from deeper ones)
+    parent_locals_count: usize,
+    /// Direct parent's upvalues (shared)
+    enclosing_upvalues: Option<Rc<RefCell<Vec<Upvalue>>>>,
+    /// Grandparent's upvalues (shared, for 3-level transitive capture)
+    grandparent_upvalues: Option<Rc<RefCell<Vec<Upvalue>>>>,
     scope_depth: u8,
     loop_stack: Vec<LoopContext>,
 }
@@ -111,9 +116,11 @@ impl Compiler {
         Self {
             function: FunctionProto::new(name),
             locals: Vec::new(),
-            upvalues: Vec::new(),
+            upvalues: Rc::new(RefCell::new(Vec::new())),
             enclosing_locals: None,
+            parent_locals_count: 0,
             enclosing_upvalues: None,
+            grandparent_upvalues: None,
             scope_depth: 0,
             loop_stack: Vec::new(),
         }
@@ -685,8 +692,15 @@ impl Compiler {
     /// Compile a nested function (lambda or function declaration) and return its proto.
     fn compile_nested_function(&mut self, name: Option<String>, params: &[String], body: &[Stmt]) -> Result<FunctionProto, RuntimeError> {
         let mut child = Compiler::new(name);
-        child.enclosing_locals = Some(self.locals.clone());
-        child.enclosing_upvalues = Some(self.upvalues.clone());
+        // Build merged enclosing locals: our locals + our enclosing chain
+        let mut all_locals = self.locals.clone();
+        if let Some(ref enc) = self.enclosing_locals {
+            all_locals.extend(enc.iter().cloned());
+        }
+        child.enclosing_locals = Some(all_locals);
+        child.parent_locals_count = self.locals.len();
+        child.enclosing_upvalues = Some(Rc::clone(&self.upvalues));
+        child.grandparent_upvalues = self.enclosing_upvalues.clone();
 
         // Add params as locals without StoreLocal — they're already on stack from Call
         for param in params {
@@ -709,7 +723,7 @@ impl Compiler {
         child.emit_op(OpCode::Return);
 
         // Mark captured locals in the parent compiler
-        for uv in &child.upvalues {
+        for uv in child.upvalues.borrow().iter() {
             if uv.is_local {
                 if let Some(local) = self.locals.get_mut(uv.index as usize) {
                     local.is_captured = true;
@@ -717,8 +731,10 @@ impl Compiler {
             }
         }
 
-        child.function.upvalue_count = child.upvalues.len() as u8;
-        child.function.upvalues = child.upvalues.iter().map(|uv| (uv.is_local, uv.index)).collect();
+        let child_upvalues = child.upvalues.borrow();
+        child.function.upvalue_count = child_upvalues.len() as u8;
+        child.function.upvalues = child_upvalues.iter().map(|uv| (uv.is_local, uv.index)).collect();
+        drop(child_upvalues);
         Ok(child.function)
     }
 
@@ -771,27 +787,99 @@ impl Compiler {
     /// Returns the upvalue index in this function's upvalues list.
     fn resolve_upvalue(&mut self, name: &str) -> Option<u8> {
         // Check if already captured
-        for (j, uv) in self.upvalues.iter().enumerate() {
+        for (j, uv) in self.upvalues.borrow().iter().enumerate() {
             if uv.name == name {
                 return Some(j as u8);
             }
         }
-        // Check enclosing function's locals
+        // Check all enclosing locals (merged chain: parent, grandparent, ...)
         if let Some(ref enclosing_locals) = self.enclosing_locals {
             for (i, local) in enclosing_locals.iter().enumerate().rev() {
                 if local.name == name && local.depth > 0 {
-                    let idx = self.upvalues.len() as u8;
-                    self.upvalues.push(Upvalue { index: i as u8, is_local: true, name: name.to_string() });
+                    if i < self.parent_locals_count {
+                        // Found in direct parent's locals → capture as upvalue
+                        let idx = self.upvalues.borrow().len() as u8;
+                        self.upvalues.borrow_mut().push(Upvalue {
+                            index: i as u8, is_local: true, name: name.to_string()
+                        });
+                        return Some(idx);
+                    } else {
+                        // Found deeper than parent. Need to create upvalue chain.
+                        if let Some(ref enc_upvalues) = self.enclosing_upvalues {
+                            // Check if parent already has this upvalue
+                            let parent_uv_idx = {
+                                let enc = enc_upvalues.borrow();
+                                enc.iter().position(|uv| uv.name == name).map(|p| p as u8)
+                            };
+                            let parent_idx = match parent_uv_idx {
+                                Some(idx) => idx,
+                                None => {
+                                    // Add upvalue to parent's list.
+                                    // The parent sees this variable at a certain index in its own
+                                    // enclosing_locals. That index is (i - self.parent_locals_count)
+                                    // in the merged list, which corresponds to the same variable
+                                    // in the parent's enclosing_locals.
+                                    let enc_idx = enc_upvalues.borrow().len() as u8;
+
+                                    // If we have grandparent_upvalues, the parent's upvalue
+                                    // should be transitive (is_local=false), and we need to
+                                    // ensure grandparent has it too.
+                                    if let Some(ref gp_upvalues) = self.grandparent_upvalues {
+                                        // Ensure grandparent has the upvalue first
+                                        let gp_idx = {
+                                            let gp = gp_upvalues.borrow();
+                                            gp.iter().position(|uv| uv.name == name).map(|p| p as u8)
+                                        };
+                                        let gp_idx = match gp_idx {
+                                            Some(idx) => idx,
+                                            None => {
+                                                let idx = gp_upvalues.borrow().len() as u8;
+                                                // Grandparent captures this as a local
+                                                // (it's directly in the grandparent's enclosing scope)
+                                                gp_upvalues.borrow_mut().push(Upvalue {
+                                                    index: (i - self.parent_locals_count) as u8,
+                                                    is_local: true,
+                                                    name: name.to_string()
+                                                });
+                                                idx
+                                            }
+                                        };
+                                        // Parent's upvalue is transitive through grandparent
+                                        enc_upvalues.borrow_mut().push(Upvalue {
+                                            index: gp_idx,
+                                            is_local: false,
+                                            name: name.to_string()
+                                        });
+                                    } else {
+                                        // No grandparent — parent captures directly as local
+                                        enc_upvalues.borrow_mut().push(Upvalue {
+                                            index: (i - self.parent_locals_count) as u8,
+                                            is_local: true,
+                                            name: name.to_string()
+                                        });
+                                    }
+                                    enc_idx
+                                }
+                            };
+                            // Add transitive upvalue in self pointing to parent's
+                            let idx = self.upvalues.borrow().len() as u8;
+                            self.upvalues.borrow_mut().push(Upvalue {
+                                index: parent_idx, is_local: false, name: name.to_string()
+                            });
                             return Some(idx);
                         }
                     }
                 }
-        // Check enclosing function's upvalues (transitive)
+            }
+        }
+        // Check parent's upvalues (transitive closure over 2 levels)
         if let Some(ref enclosing_upvalues) = self.enclosing_upvalues {
-            for uv in enclosing_upvalues.iter().rev() {
+            for uv in enclosing_upvalues.borrow().iter().rev() {
                 if uv.name == name {
-                    let idx = self.upvalues.len() as u8;
-                    self.upvalues.push(Upvalue { index: uv.index, is_local: false, name: name.to_string() });
+                    let idx = self.upvalues.borrow().len() as u8;
+                    self.upvalues.borrow_mut().push(Upvalue {
+                        index: uv.index, is_local: false, name: name.to_string()
+                    });
                     return Some(idx);
                 }
             }