update

Author: CppCXY

crates/luars/src/compiler/expr.rs

@@ -3333,6 +3333,7 @@ pub fn compile_closure_expr_to(
                 register: 0,
                 is_const: false,
                 is_to_be_closed: false,
+                needs_close: false,
             });
         func_compiler.chunk.locals.push("self".to_string());
         param_offset = 1;
@@ -3367,6 +3368,7 @@ pub fn compile_closure_expr_to(
                 register: reg_index,
                 is_const: false,
                 is_to_be_closed: false,
+                needs_close: false,
             });
         func_compiler.chunk.locals.push(param_name);
         regular_param_count += 1;

crates/luars/src/compiler/helpers.rs

@@ -199,6 +199,7 @@ pub fn add_local_with_attrs(
         register,
         is_const,
         is_to_be_closed,
+        needs_close: false,
     };
     c.scope_chain.borrow_mut().locals.push(local);
 
@@ -271,9 +272,11 @@ pub fn resolve_upvalue_from_chain(c: &mut Compiler, name: &str) -> Option<usize>
 fn resolve_in_parent_scope(scope: &Rc<RefCell<ScopeChain>>, name: &str) -> Option<(bool, u32)> {
     // First, search in this scope's locals
     {
-        let scope_ref = scope.borrow();
-        if let Some(local) = scope_ref.locals.iter().rev().find(|l| l.name == name) {
+        let mut scope_ref = scope.borrow_mut();
+        if let Some(local) = scope_ref.locals.iter_mut().rev().find(|l| l.name == name) {
             let register = local.register;
+            // Mark this local as captured by a closure - needs CLOSE on scope exit
+            local.needs_close = true;
             // Found as local - return (true, register)
             return Some((true, register));
         }
@@ -568,8 +571,15 @@ pub fn try_expr_as_constant(c: &mut Compiler, expr: &emmylua_parser::LuaExpr) ->
 
 /// Begin a new loop (for break statement support)
 pub fn begin_loop(c: &mut Compiler) {
+    begin_loop_with_register(c, c.freereg);
+}
+
+/// Begin a new loop with a specific first register for CLOSE
+pub fn begin_loop_with_register(c: &mut Compiler, first_reg: u32) {
     c.loop_stack.push(super::LoopInfo {
         break_jumps: Vec::new(),
+        scope_depth: c.scope_depth,
+        first_local_register: first_reg,
     });
 }
 
@@ -589,6 +599,35 @@ pub fn emit_break(c: &mut Compiler) -> Result<(), String> {
         return Err("break statement outside loop".to_string());
     }
 
+    // Before breaking, check if we need to close any captured upvalues
+    // This is needed when break jumps past the CLOSE instruction that would
+    // normally be executed at the end of each iteration
+    let loop_info = c.loop_stack.last().unwrap();
+    let loop_scope_depth = loop_info.scope_depth;
+    let first_reg = loop_info.first_local_register;
+    
+    // Find minimum register of captured locals in the loop scope
+    let mut min_close_reg: Option<u32> = None;
+    {
+        let scope = c.scope_chain.borrow();
+        for local in scope.locals.iter().rev() {
+            if local.depth < loop_scope_depth {
+                break; // Only check loop scope and nested scopes
+            }
+            if local.needs_close && local.register >= first_reg {
+                min_close_reg = Some(match min_close_reg {
+                    None => local.register,
+                    Some(min_reg) => min_reg.min(local.register),
+                });
+            }
+        }
+    }
+    
+    // Emit CLOSE if needed
+    if let Some(reg) = min_close_reg {
+        emit(c, Instruction::encode_abc(OpCode::Close, reg, 0, 0));
+    }
+
     let jump_pos = emit_jump(c, OpCode::Jmp);
     c.loop_stack.last_mut().unwrap().break_jumps.push(jump_pos);
     Ok(())

crates/luars/src/compiler/mod.rs

@@ -83,11 +83,14 @@ pub(crate) struct Local {
     pub register: u32,
     pub is_const: bool,        // <const> attribute
     pub is_to_be_closed: bool, // <close> attribute
+    pub needs_close: bool,     // True if captured by a closure (needs CLOSE on scope exit)
 }
 
 /// Loop information for break statements
 pub(crate) struct LoopInfo {
-    pub break_jumps: Vec<usize>, // Positions of break statements to patch
+    pub break_jumps: Vec<usize>,   // Positions of break statements to patch
+    pub scope_depth: usize,        // Scope depth at loop start
+    pub first_local_register: u32, // First register of loop-local variables (for CLOSE on break)
 }
 
 /// Label definition

crates/luars/src/compiler/stmt.rs

@@ -1223,9 +1223,12 @@ fn compile_for_stat(c: &mut Compiler, stat: &LuaForStat) -> Result<(), String> {
     // Begin new scope for loop body
     begin_scope(c);
 
+    // Begin loop with var_reg as first register, so break can close it
+    // Using var_reg instead of c.freereg because var_reg is the loop variable
+    begin_loop_with_register(c, var_reg);
+
     // The loop variable is at R(base+3)
     add_local(c, var_name, var_reg);
-    begin_loop(c);
 
     // Loop body starts here
     let loop_body_start = c.chunk.code.len();
@@ -1235,7 +1238,29 @@ fn compile_for_stat(c: &mut Compiler, stat: &LuaForStat) -> Result<(), String> {
         compile_block(c, &body)?;
     }
 
-    // FORLOOP comes AFTER the body
+    // Before FORLOOP, emit CLOSE if any local in the loop body was captured
+    // Find the minimum register of captured locals in the current scope (loop body)
+    {
+        let scope = c.scope_chain.borrow();
+        let mut min_close_reg: Option<u32> = None;
+        for local in scope.locals.iter().rev() {
+            if local.depth < c.scope_depth {
+                break; // Only check current scope (loop body)
+            }
+            if local.depth == c.scope_depth && local.needs_close {
+                min_close_reg = Some(match min_close_reg {
+                    None => local.register,
+                    Some(min_reg) => min_reg.min(local.register),
+                });
+            }
+        }
+        drop(scope);
+        if let Some(reg) = min_close_reg {
+            emit(c, Instruction::encode_abc(OpCode::Close, reg, 0, 0));
+        }
+    }
+
+    // FORLOOP comes AFTER the body (and CLOSE if needed)
     let forloop_pc = c.chunk.code.len();
     // Emit FORLOOP: increments index, checks condition, copies to var, jumps back to body
     // Bx is the backward jump distance. Since PC is incremented before dispatch,
@@ -1516,6 +1541,7 @@ fn compile_local_function_stat(
         register: func_reg,
         is_const: false,
         is_to_be_closed: false,
+        needs_close: false,
     });
     c.chunk.locals.push(func_name.clone());