aster-core/src/vm/compiler.rs

@@ -84,6 +84,7 @@ struct Local {
 struct Upvalue {
     index: u8,
     is_local: bool,  // true = captured from enclosing fn's local; false = from upvalue
+    name: String,    // variable name (for transitive upvalue resolution)
 }
 
 struct LoopContext {
@@ -97,9 +98,9 @@ pub struct Compiler {
     function: FunctionProto,
     locals: Vec<Local>,
     upvalues: Vec<Upvalue>,
-    /// Snapshot of enclosing compiler's locals (for upvalue resolution)
+    /// Direct parent's locals (for upvalue resolution)
     enclosing_locals: Option<Vec<Local>>,
-    /// Snapshot of enclosing compiler's upvalues (for transitive upvalues)
+    /// Direct parent's upvalues (for transitive upvalue resolution)
     enclosing_upvalues: Option<Vec<Upvalue>>,
     scope_depth: u8,
     loop_stack: Vec<LoopContext>,
@@ -303,13 +304,39 @@ impl Compiler {
     }
 
     fn compile_for_in(&mut self, var_name: &str, iterable: &Expr, body: &Stmt) -> Result<(), RuntimeError> {
-        // Evaluate iterable, set up iterator
-        self.compile_expr(iterable)?;
-        emit_op(&mut self.function.code, OpCode::ForInInit);
-        let exit_jump = emit_i16_placeholder(&mut self.function.code, OpCode::ForInNext);
-
-        // New scope for the loop variable
+        // 1. New scope for hidden iterator locals (items, idx)
         self.begin_scope();
+
+        // 2. Evaluate iterable and compute items array
+        self.compile_expr(iterable)?;
+        emit_op(&mut self.function.code, OpCode::ForInInit); // pops iterable, pushes items[]
+
+        // 3. Store items array in a hidden local (peek, no Pop — cleaned by end_scope)
+        let items_slot = self.locals.len() as u8;
+        self.locals.push(Local {
+            name: format!("__iter_items_{}", items_slot),
+            depth: self.scope_depth,
+            is_captured: false,
+            is_const: false,
+        });
+        emit_u8(&mut self.function.code, OpCode::StoreLocal, items_slot);
+
+        // 4. Initialize index = 0 in hidden local
+        let idx_slot = self.locals.len() as u8;
+        self.locals.push(Local {
+            name: format!("__iter_idx_{}", idx_slot),
+            depth: self.scope_depth,
+            is_captured: false,
+            is_const: false,
+        });
+        let zero_idx = self.add_constant(Value::Number(0.0));
+        emit_u16(&mut self.function.code, OpCode::LoadConst, zero_idx);
+        emit_u8(&mut self.function.code, OpCode::StoreLocal, idx_slot);
+
+        // 5. ForInNext: reads items[idx_slot], idx[idx_slot], pushes element
+        let forin_loc = self.emit_forin_next(items_slot, idx_slot);
+
+        // 6. Store loop variable in a local (peek, value stays on stack)
         let loop_var_slot = self.locals.len() as u8;
         self.locals.push(Local {
             name: var_name.to_string(),
@@ -318,13 +345,10 @@ impl Compiler {
             is_const: false,
         });
         emit_u8(&mut self.function.code, OpCode::StoreLocal, loop_var_slot);
-        // Pop the ForInNext value from stack top (it's stored in the local)
-        self.emit_op(OpCode::Pop);
-
-        let start_ip = self.function.code.len();
 
+        // 7. Loop body (can resolve var_name to loop_var_slot)
         self.loop_stack.push(LoopContext {
-            start_ip,
+            start_ip: forin_loc,
             break_patches: Vec::new(),
             continue_patches: Vec::new(),
             scope_depth: self.scope_depth,
@@ -332,28 +356,31 @@ impl Compiler {
 
         self.compile_stmt(body)?;
 
-        // Patch continue → loop back
+        // Patch continue
         let continue_ip = self.function.code.len();
         let loop_ctx = self.loop_stack.pop().unwrap();
         for patch in loop_ctx.continue_patches {
             self.patch_jump_to(patch, continue_ip);
         }
 
-        // Jump back to ForInNext
-        self.emit_loop_jump(exit_jump);
-        let exit_ip = self.function.code.len();
-        self.patch_jump_to(exit_jump, exit_ip);
+        // 8. Pop loop var element and remove from tracking
+        self.emit_op(OpCode::Pop);
+        self.locals.pop(); // loop_var_slot
 
+        // 9. Jump back to ForInNext (next iteration re-adds loop var via StoreLocal)
+        self.emit_loop_jump(forin_loc);
+
+        // 10. Exit target: patch ForInNext exit and break jumps
+        let exit_ip = self.function.code.len();
+        self.patch_forin_jump(forin_loc, exit_ip);
         for patch in loop_ctx.break_patches {
-            self.patch_jump(patch);
+            self.patch_jump_to(patch, exit_ip);
         }
 
-        // Pop the loop variable
-        self.emit_op(OpCode::Pop);
-        // Pop the iterator state (2 values: iterable ref + index)
-        self.emit_op(OpCode::Pop);
+        // 11. Pop loop var (ForInNext pushes Nil on exit; break leaves element on stack)
         self.emit_op(OpCode::Pop);
 
+        // 12. End scope: pops items_slot + idx_slot values (left by StoreLocal peek)
         self.end_scope();
         Ok(())
     }
@@ -658,7 +685,6 @@ 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);
-        // Pass enclosing state for upvalue resolution
         child.enclosing_locals = Some(self.locals.clone());
         child.enclosing_upvalues = Some(self.upvalues.clone());
 
@@ -744,22 +770,32 @@ impl Compiler {
     /// Try to resolve a variable as an upvalue from enclosing functions.
     /// Returns the upvalue index in this function's upvalues list.
     fn resolve_upvalue(&mut self, name: &str) -> Option<u8> {
-        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 {
-                    // Check if we already captured this upvalue
+        // Check if already captured
         for (j, uv) in self.upvalues.iter().enumerate() {
-                        if uv.is_local && uv.index == i as u8 {
+            if uv.name == name {
                 return Some(j as u8);
             }
         }
+        // Check enclosing function's locals
+        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 });
+                    self.upvalues.push(Upvalue { index: i as u8, is_local: true, name: name.to_string() });
+                    return Some(idx);
+                }
+            }
+        }
+        // Check enclosing function's upvalues (transitive)
+        if let Some(ref enclosing_upvalues) = self.enclosing_upvalues {
+            for uv in enclosing_upvalues.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() });
                     return Some(idx);
                 }
             }
         }
-        // Transitive upvalues (from enclosing's upvalues) not yet implemented
         None
     }
 
@@ -788,6 +824,17 @@ impl Compiler {
         emit_i16(&mut self.function.code, OpCode::Jump, offset as i16);
     }
 
+    fn emit_forin_next(&mut self, items_slot: u8, idx_slot: u8) -> usize {
+        let loc = self.function.code.len();
+        let code = &mut self.function.code;
+        code.push(OpCode::ForInNext as u8);
+        code.push(items_slot);
+        code.push(idx_slot);
+        code.push(0xFF); // placeholder offset low
+        code.push(0x7F); // placeholder offset high
+        loc
+    }
+
     fn emit_closure(&mut self, proto_idx: u16, upvalues: &[(bool, u8)]) {
         let code = &mut self.function.code;
         code.push(OpCode::Closure as u8);
@@ -815,6 +862,18 @@ impl Compiler {
         code[jump_loc + 2] = ((offset >> 8) & 0xFF) as u8;
     }
 
+    /// Patch ForInNext's exit offset (at jump_loc + 3, +4)
+    fn patch_forin_jump(&mut self, forin_loc: usize, target: usize) {
+        let offset = target as isize - forin_loc as isize;
+        let code = &mut self.function.code;
+        code[forin_loc + 3] = (offset & 0xFF) as u8;
+        code[forin_loc + 4] = ((offset >> 8) & 0xFF) as u8;
+    }
+
+    fn add_constant(&mut self, val: Value) -> u16 {
+        self.function.add_constant(val)
+    }
+
     // ========================================================================
     // Constant pool helpers
     // ========================================================================

aster-core/src/vm/vm.rs

@@ -511,32 +511,41 @@ impl Vm {
                 // --- For-in ---
                 OpCode::ForInInit => {
                     let iterable = self.stack.pop().unwrap();
-                    // Push iterator state: (collection, index)
-                    let iter = Value::Number(0.0);
-                    self.stack.push(iterable);
-                    self.stack.push(iter);
+                    let items = self.for_in_items(&iterable);
+                    self.stack.push(Value::Array(Rc::new(RefCell::new(items))));
                     self.advance_ip(SIZE_OP);
                 }
                 OpCode::ForInNext => {
-                    let exit_offset = read_i16(code, ip) as isize;
-                    let exit_ip = ((ip as isize) + exit_offset) as usize;
-                    let iter_idx = self.stack.pop().unwrap(); // current index
-                    let collection = self.stack.pop().unwrap(); // the iterable
+                    // Encoding: opcode(1) + items_slot(1) + idx_slot(1) + exit_offset(2) = 5
+                    let items_slot = code[ip + 1] as usize;
+                    let idx_slot = code[ip + 2] as usize;
+                    let exit_offset = ((code[ip + 3] as u16) | ((code[ip + 4] as u16) << 8)) as i16;
+                    let exit_ip = ((ip as isize) + exit_offset as isize) as usize;
+                    let base = self.frames.last().unwrap().stack_base;
+
+                    // Read index from local slot
+                    let idx_val = self.stack[base + idx_slot].clone();
+                    let idx = self.as_usize(&idx_val, "for-in index")?;
+                    let items_val = self.stack[base + items_slot].clone();
+
+                    let items = match &items_val {
+                        Value::Array(arr) => arr.borrow(),
+                        _ => return Err(RuntimeError::RuntimeError {
+                            message: "ForInNext: items is not an array".into(),
+                            token: None,
+                        }),
+                    };
 
-                    let idx = self.as_number(&iter_idx)? as usize;
-                    let items = self.for_in_items(&collection);
                     if idx >= items.len() {
-                        // Done iterating — push back state and jump to exit
-                        self.stack.push(collection);
-                        self.stack.push(iter_idx);
+                        // Done: push Nil (keeps stack balanced for Pop at exit)
+                        self.stack.push(Value::Nil);
                         self.advance_ip_to(exit_ip);
                     } else {
-                        // Push back incremented state
-                        self.stack.push(collection);
-                        self.stack.push(Value::Number((idx + 1) as f64));
-                        // Push the current value for the loop body
+                        // Push current element for loop body
                         self.stack.push(items[idx].clone());
-                        self.advance_ip(SIZE_U16);
+                        // Increment index in local slot
+                        self.stack[base + idx_slot] = Value::Number((idx + 1) as f64);
+                        self.advance_ip(5); // SIZE_FORIN_NEXT = 5
                     }
                 }