Feat: Implement list set dict and generator comprehensions via drain-reinject with SSA slot rewriting and append opcodes

dylan-sutton-chavez · dylan-sutton-chavez · commit b82c452b03e8 · 2026-04-20T19:27:18.000-06:00
diff --git a/compiler/src/modules/parser/expr.rs b/compiler/src/modules/parser/expr.rs
@@ -1,7 +1,7 @@
 // parser/expr.rs
 
 use super::Parser;
-use super::types::{OpCode, Value, MAX_EXPR_DEPTH};
+use super::types::{OpCode, Value, MAX_EXPR_DEPTH, Instruction};
 use super::types::parse_string;
 use crate::modules::lexer::{Token, TokenType};
 use alloc::{string::ToString, vec::Vec, vec, format, string::String};
@@ -183,22 +183,31 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
             TokenType::Lbrace => self.brace_literal(),
             TokenType::Lsqb => self.list_literal(),
             TokenType::Lpar => {
-                if matches!(self.peek(), Some(TokenType::Rpar)) { self.advance(); self.chunk.emit(OpCode::BuildTuple, 0); } else {
-                self.expr();
-                if matches!(self.peek(), Some(TokenType::For)) {
-                    self.comprehension(OpCode::GenExpr);
-                } else if self.eat_if(TokenType::Comma) {
-                    let mut count = 1u16;
-                    while !matches!(self.peek(), Some(TokenType::Rpar) | None) {
-                        self.expr();
-                        count += 1;
-                        if !self.eat_if(TokenType::Comma) { break; }
-                    }
-                    self.eat(TokenType::Rpar);
-                    self.chunk.emit(OpCode::BuildTuple, count);
+                if matches!(self.peek(), Some(TokenType::Rpar)) {
+                    self.advance();
+                    self.chunk.emit(OpCode::BuildTuple, 0);
                 } else {
-                    self.eat(TokenType::Rpar);
-                }}
+                    let elem_start = self.chunk.instructions.len();
+                    self.expr();
+                    if matches!(self.peek(), Some(TokenType::For)) {
+                        let versions_before = self.ssa_versions.clone();
+                        let elem_ins: Vec<Instruction> = self.chunk.instructions.drain(elem_start..).collect();
+                        self.chunk.emit(OpCode::BuildList, 0);
+                        self.comprehension_loop(&[elem_ins], OpCode::ListAppend, &versions_before);
+                        self.advance(); // Rpar
+                    } else if self.eat_if(TokenType::Comma) {
+                        let mut count = 1u16;
+                        while !matches!(self.peek(), Some(TokenType::Rpar) | None) {
+                            self.expr();
+                            count += 1;
+                            if !self.eat_if(TokenType::Comma) { break; }
+                        }
+                        self.eat(TokenType::Rpar);
+                        self.chunk.emit(OpCode::BuildTuple, count);
+                    } else {
+                        self.eat(TokenType::Rpar);
+                    }
+                }
             }
             TokenType::Lambda => self.parse_lambda(),
             _ => {
diff --git a/compiler/src/modules/parser/literals.rs b/compiler/src/modules/parser/literals.rs
@@ -3,9 +3,10 @@
 use super::Parser;
 use super::types::builtin;
 
-use super::types::{OpCode, Value, SSAChunk};
+use super::types::{OpCode, Value, SSAChunk, Instruction};
 use crate::modules::lexer::{Token, TokenType};
 use alloc::{string::{String, ToString}, vec::Vec, format};
+use hashbrown::HashMap;
 
 impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
 
@@ -20,19 +21,24 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
             self.chunk.emit(OpCode::BuildDict, 0);
             return;
         }
+        let key_start = self.chunk.instructions.len();
         self.expr();
         match self.peek() {
             Some(TokenType::Colon) => {
                 self.advance();
+                let val_start = self.chunk.instructions.len();
                 self.expr();
                 if matches!(self.peek(), Some(TokenType::For)) {
-                    self.comprehension(OpCode::DictComp);
+                    let versions_before = self.ssa_versions.clone();
+                    let val_ins: Vec<Instruction> = self.chunk.instructions.drain(val_start..).collect();
+                    let key_ins: Vec<Instruction> = self.chunk.instructions.drain(key_start..).collect();
+                    self.chunk.emit(OpCode::BuildDict, 0);
+                    self.comprehension_loop(&[key_ins, val_ins], OpCode::MapAdd, &versions_before);
+                    self.advance(); // Rbrace
                 } else {
                     let mut pairs = 1u16;
                     while self.eat_if(TokenType::Comma) {
-                        if matches!(self.peek(), Some(TokenType::Rbrace)) {
-                            break;
-                        }
+                        if matches!(self.peek(), Some(TokenType::Rbrace)) { break; }
                         self.expr();
                         self.eat(TokenType::Colon);
                         self.expr();
@@ -43,14 +49,16 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
                 }
             }
             Some(TokenType::For) => {
-                self.comprehension(OpCode::SetComp);
+                let versions_before = self.ssa_versions.clone();
+                let elem_ins: Vec<Instruction> = self.chunk.instructions.drain(key_start..).collect();
+                self.chunk.emit(OpCode::BuildSet, 0);
+                self.comprehension_loop(&[elem_ins], OpCode::SetAdd, &versions_before);
+                self.advance(); // Rbrace
             }
             _ => {
                 let mut count = 1u16;
                 while self.eat_if(TokenType::Comma) {
-                    if matches!(self.peek(), Some(TokenType::Rbrace)) {
-                        break;
-                    }
+                    if matches!(self.peek(), Some(TokenType::Rbrace)) { break; }
                     self.expr();
                     count += 1;
                 }
@@ -71,15 +79,18 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
             self.chunk.emit(OpCode::BuildList, 0);
             return;
         }
+        let elem_start = self.chunk.instructions.len();
         self.expr();
         if matches!(self.peek(), Some(TokenType::For)) {
-            self.comprehension(OpCode::ListComp);
+            let versions_before = self.ssa_versions.clone();
+            let elem_ins: Vec<Instruction> = self.chunk.instructions.drain(elem_start..).collect();
+            self.chunk.emit(OpCode::BuildList, 0);
+            self.comprehension_loop(&[elem_ins], OpCode::ListAppend, &versions_before);
+            self.advance(); // Rsqb
         } else {
             let mut count = 1u16;
             while self.eat_if(TokenType::Comma) {
-                if matches!(self.peek(), Some(TokenType::Rsqb)) {
-                    break;
-                }
+                if matches!(self.peek(), Some(TokenType::Rsqb)) { break; }
                 self.expr();
                 count += 1;
             }
@@ -89,25 +100,22 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
     }
 
     /*
-    Comprehension Handler
-        Generates SSA bytecode for list/set/dict/generator comprehensions.
+    Comprehension Loop
+        Builds for/if scaffolding and reinjects element body with SSA slot rewriting.
     */
 
-    pub(super) fn comprehension(&mut self, op: OpCode) {
-        let mut loop_starts = Vec::new();
-        let mut for_iters = Vec::new();
+    pub(super) fn comprehension_loop(&mut self, elem_bodies: &[Vec<Instruction>], append_op: OpCode, versions_before: &HashMap<String, u32>) {
+        let mut loop_starts: Vec<u16> = Vec::new();
+        let mut for_iters: Vec<usize> = Vec::new();
+        let mut all_vars: Vec<String> = Vec::new();
 
         while self.eat_if(TokenType::For) {
-            let mut vars = Vec::new();
+            let mut vars: Vec<String> = Vec::new();
             loop {
                 let t = self.advance();
                 vars.push(self.lexeme(&t).to_string());
-                if !self.eat_if(TokenType::Comma) {
-                    break;
-                }
-                if matches!(self.peek(), Some(TokenType::In)) {
-                    break;
-                }
+                if !self.eat_if(TokenType::Comma) { break; }
+                if matches!(self.peek(), Some(TokenType::In)) { break; }
             }
 
             self.eat(TokenType::In);
@@ -119,42 +127,54 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
             let fi = self.chunk.instructions.len() - 1;
 
             if vars.len() == 1 {
-                let ver = self.increment_version(&vars[0]);
-                let mut buf = [0u8; 128];
-                let idx = self.chunk.push_name(Self::ssa_name(&vars[0], ver, &mut buf));
-                self.chunk.emit(OpCode::StoreName, idx);
+                self.store_name(vars[0].clone());
             } else {
                 self.chunk.emit(OpCode::UnpackSequence, vars.len() as u16);
                 for var in vars.iter().rev() {
                     self.store_name(var.clone());
                 }
             }
+            for v in &vars { all_vars.push(v.clone()); }
 
             while self.eat_if(TokenType::If) {
                 self.expr_bp(1);
-                self.chunk.emit(OpCode::JumpIfFalse, ls);
+                self.chunk.emit(OpCode::JumpIfFalse, ls); // fail -> next iteration
             }
 
             loop_starts.push(ls);
             for_iters.push(fi);
         }
 
-        let n = for_iters.len();
-        let mut jump_positions = Vec::new();
-        for &ls in loop_starts.iter().rev() {
-            self.chunk.emit(OpCode::Jump, ls);
-            jump_positions.push(self.chunk.instructions.len() - 1);
+        // Map pre-loop slots to current versions, skipping non-existent element references.
+        let mut var_map: HashMap<u16, u16> = HashMap::new();
+        for var in &all_vars {
+            let old_ver = versions_before.get(var).copied().unwrap_or(0);
+            let new_ver = self.current_version(var);
+            if old_ver == new_ver { continue; }
+            let old_name = format!("{}_{}", var, old_ver);
+            let Some(&old_slot) = self.chunk.name_index.get(old_name.as_str()) else { continue };
+            let mut nb = [0u8; 128];
+            let new_slot = self.chunk.push_name(Self::ssa_name(var, new_ver, &mut nb));
+            var_map.insert(old_slot, new_slot);
         }
 
-        for i in 1..n {
-            let target = jump_positions[n - i] as u16;
-            self.chunk.instructions[for_iters[i]].operand = target; // back-patch inner ForIter now that all loop headers are known
+        for body in elem_bodies {
+            for ins in body {
+                let operand = if matches!(ins.opcode, OpCode::LoadName | OpCode::StoreName) {
+                    var_map.get(&ins.operand).copied().unwrap_or(ins.operand)
+                } else {
+                    ins.operand
+                };
+                self.chunk.instructions.push(Instruction { opcode: ins.opcode, operand });
+            }
         }
+        self.chunk.emit(append_op, 0);
 
-        self.patch(for_iters[0]);
-
-        self.advance();
-        self.chunk.emit(op, 0);
+        // Close loops innermost-first: Jump back to header, patch matching ForIter to the point past it.
+        for i in (0..for_iters.len()).rev() {
+            self.chunk.emit(OpCode::Jump, loop_starts[i]);
+            self.patch(for_iters[i]);
+        }
     }
 
     /*
diff --git a/compiler/src/modules/parser/types.rs b/compiler/src/modules/parser/types.rs
@@ -21,8 +21,8 @@ pub enum OpCode {
     GetIter, ForIter, GetItem, Mod, Pow, FloorDiv, LoadTrue, LoadFalse, LoadNone, LoadAttr, StoreAttr, 
     BuildSlice, MakeClass, SetupExcept, PopExcept, Raise, Import, ImportFrom, BitAnd, BitOr, BitXor, 
     BitNot, Shl, Shr, In, NotIn, Is, IsNot, UnpackSequence, BuildTuple, SetupWith, ExitWith, Yield, 
-    Del, Assert, Global, Nonlocal, UnpackArgs, ListComp, SetComp, DictComp, BuildSet, RaiseFrom, 
-    UnpackEx, LoadEllipsis, GenExpr, Await, MakeCoroutine, YieldFrom, TypeAlias, StoreItem, Dup2
+    Del, Assert, Global, Nonlocal, UnpackArgs, ListAppend, SetAdd, MapAdd, BuildSet, RaiseFrom, 
+    UnpackEx, LoadEllipsis, Await, MakeCoroutine, YieldFrom, TypeAlias, StoreItem, Dup2
 }
 
 /*
diff --git a/compiler/src/modules/vm/mod.rs b/compiler/src/modules/vm/mod.rs
@@ -513,6 +513,40 @@ impl<'a> VM<'a> {
                     }
                 }
 
+                // Comprehension append
+
+                OpCode::ListAppend => {
+                    let v = self.pop()?;
+                    let acc = *self.stack.last().ok_or_else(|| VmErr::Runtime("stack underflow"))?;
+                    if !acc.is_heap() { return Err(VmErr::Runtime("list accumulator corrupted")); }
+                    match self.heap.get(acc) {
+                        HeapObj::List(rc) => rc.borrow_mut().push(v),
+                        _ => return Err(VmErr::Runtime("list accumulator corrupted")),
+                    }
+                }
+                OpCode::SetAdd => {
+                    let v = self.pop()?;
+                    let acc = *self.stack.last().ok_or_else(|| VmErr::Runtime("stack underflow"))?;
+                    if !acc.is_heap() { return Err(VmErr::Runtime("set accumulator corrupted")); }
+                    let already = match self.heap.get(acc) {
+                        HeapObj::Set(rc) => rc.borrow().iter().any(|&x| eq_vals_with_heap(x, v, &self.heap)),
+                        _ => return Err(VmErr::Runtime("set accumulator corrupted")),
+                    };
+                    if !already {
+                        if let HeapObj::Set(rc) = self.heap.get(acc) { rc.borrow_mut().push(v); }
+                    }
+                }
+                OpCode::MapAdd => {
+                    let value = self.pop()?;
+                    let key = self.pop()?;
+                    let acc = *self.stack.last().ok_or_else(|| VmErr::Runtime("stack underflow"))?;
+                    if !acc.is_heap() { return Err(VmErr::Runtime("dict accumulator corrupted")); }
+                    match self.heap.get(acc) {
+                        HeapObj::Dict(rc) => { rc.borrow_mut().insert(key, value); }
+                        _ => return Err(VmErr::Runtime("dict accumulator corrupted")),
+                    }
+                }
+
                 // SSA Phi
 
                 OpCode::Phi => { Self::exec_phi(op, rip, &phi_map, slots, prev_slots, &chunk.phi_sources); }
@@ -656,8 +690,6 @@ impl<'a> VM<'a> {
                 OpCode::UnpackArgs => { return Err(VmErr::Runtime("*args/**kwargs not yet supported")); }
                 OpCode::MakeClass => { return Err(VmErr::Runtime("classes not yet supported")); }
                 OpCode::LoadAttr | OpCode::StoreAttr => { return Err(VmErr::Runtime("attribute access not yet supported")); }
-                OpCode::ListComp | OpCode::SetComp | OpCode::DictComp => { return Err(VmErr::Runtime("comprehensions not yet supported")); }
-                OpCode::GenExpr => { return Err(VmErr::Runtime("generator expressions not yet supported")); }
             }
         }
     }
diff --git a/compiler/src/wasm.rs b/compiler/src/wasm.rs
@@ -1,9 +1,7 @@
-use crate::modules::{lexer::lexer, parser::Parser, vm::{VM, Limits}};
-
 #[cfg(target_arch = "wasm32")]
 mod runtime {
     use lol_alloc::LeakingPageAllocator;
-    use super::{lexer, Parser, VM, Limits};
+    use crate::modules::{lexer::lexer, parser::Parser, vm::{VM, Limits}};
 
     #[global_allocator]
     static A: LeakingPageAllocator = LeakingPageAllocator;
diff --git a/compiler/tests/cases/parser_cases.json b/compiler/tests/cases/parser_cases.json
@@ -859,11 +859,11 @@
     "annotations": {}
   },
   {
-    "src": "x = [i for i in range(3)]", 
-    "constants": ["3"], 
-    "names": ["i_0", "i_1", "x_1"], 
-    "instructions": [["LoadName",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",7], ["StoreName",1], ["Jump",4], ["ListComp",0], ["StoreName",2], ["ReturnValue",0]], 
-    "annotations": {}
+      "src": "x = [i for i in range(3)]",
+      "constants": ["3"],
+      "names": ["i_0", "i_1", "x_1"],
+      "instructions": [["BuildList",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",9], ["StoreName",1], ["LoadName",1], ["ListAppend",0], ["Jump",4], ["StoreName",2], ["ReturnValue",0]],
+      "annotations": {}
   },
   {
     "src": "x = (y := 5)", 
@@ -1021,11 +1021,11 @@
     "annotations": {}
   },
   {
-    "src": "x = [i for i in range(3) for j in range(2)]", 
-    "constants": ["3", "2"], 
-    "names": ["i_0", "i_1", "j_1", "x_1"], 
-    "instructions": [["LoadName",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",13], ["StoreName",1], ["LoadConst",1], ["CallRange",1], ["GetIter",0], ["ForIter",12], ["StoreName",2], ["Jump",9], ["Jump",4], ["ListComp",0], ["StoreName",3], ["ReturnValue",0]], 
-    "annotations": {}
+      "src": "x = [i for i in range(3) for j in range(2)]",
+      "constants": ["3", "2"],
+      "names": ["i_0", "i_1", "j_1", "x_1"],
+      "instructions": [["BuildList",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",15], ["StoreName",1], ["LoadConst",1], ["CallRange",1], ["GetIter",0], ["ForIter",14], ["StoreName",2], ["LoadName",1], ["ListAppend",0], ["Jump",9], ["Jump",4], ["StoreName",3], ["ReturnValue",0]],
+      "annotations": {}
   },
   {
     "src": "with a, b:\n    pass", 
@@ -1063,17 +1063,17 @@
     "annotations": {}
   },
   {
-    "src": "x = [i for i in range(10) if i > 2 if i < 8]", 
-    "constants": ["10", "2", "8"], 
-    "names": ["i_0", "i_1", "x_1"], 
-    "instructions": [["LoadName",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",15], ["StoreName",1], ["LoadName",1], ["LoadConst",1], ["Gt",0], ["JumpIfFalse",4], ["LoadName",1], ["LoadConst",2], ["Lt",0], ["JumpIfFalse",4], ["Jump",4], ["ListComp",0], ["StoreName",2], ["ReturnValue",0]], 
+    "src": "x = [i for i in range(10) if i > 2 if i < 8]",
+    "constants": ["10", "2", "8"],
+    "names": ["i_0", "i_1", "x_1"],
+    "instructions": [["BuildList",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",17], ["StoreName",1], ["LoadName",1], ["LoadConst",1], ["Gt",0], ["JumpIfFalse",4], ["LoadName",1], ["LoadConst",2], ["Lt",0], ["JumpIfFalse",4], ["LoadName",1], ["ListAppend",0], ["Jump",4], ["StoreName",2], ["ReturnValue",0]],
     "annotations": {}
   },
   {
-    "src": "x = (i for i in range(3))", 
-    "constants": ["3"], 
-    "names": ["i_0", "i_1", "x_1"], 
-    "instructions": [["LoadName",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",7], ["StoreName",1], ["Jump",4], ["GenExpr",0], ["StoreName",2], ["ReturnValue",0]], 
+    "src": "x = (i for i in range(3))",
+    "constants": ["3"],
+    "names": ["i_0", "i_1", "x_1"],
+    "instructions": [["BuildList",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",9], ["StoreName",1], ["LoadName",1], ["ListAppend",0], ["Jump",4], ["StoreName",2], ["ReturnValue",0]],
     "annotations": {}
   },
   {
@@ -1541,12 +1541,12 @@
     "annotations": {}
   },
   {
-    "src": "x = {k: v for k, v in items}", 
-    "constants": [], 
-    "names": ["k_0","v_0","items_0","v_1","k_1","x_1"], 
-    "instructions": [["LoadName",0],["LoadName",1],["LoadName",2],["GetIter",0],["ForIter",9],["UnpackSequence",2],["StoreName",3],["StoreName",4],["Jump",4],["DictComp",0],["StoreName",5],["ReturnValue",0]], 
+    "src": "x = {k: v for k, v in items}",
+    "constants": [],
+    "names": ["k_0", "v_0", "items_0", "v_1", "k_1", "x_1"],
+    "instructions": [["BuildDict",0], ["LoadName",2], ["GetIter",0], ["ForIter",11], ["UnpackSequence",2], ["StoreName",3], ["StoreName",4], ["LoadName",4], ["LoadName",3], ["MapAdd",0], ["Jump",3], ["StoreName",5], ["ReturnValue",0]],
     "annotations": {}
-  },
+  }
   {
     "src": "a[0] = 1", 
     "constants": ["0","1"], 
diff --git a/compiler/tests/cases/vm_cases.json b/compiler/tests/cases/vm_cases.json
