Fix: Support generator expressions as function arguments.

Author: dylan-sutton-chavez

compiler/src/modules/parser/literals.rs

@@ -297,13 +297,29 @@ impl<'src, I: Iterator<Item = Token>> Parser<'src, I> {
                     let i = self.chunk.push_const(Value::Str(self.lexeme(&t).to_string()));
                     self.chunk.emit(OpCode::LoadConst, i);
                     self.expr();
-                    kwargs += 1; // +1 in argc below -> 2 slots per kwarg matches key+value on stack
+                    kwargs += 1;
                 } else {
+                    // Support genexpr after name arg (expr for ...)
+                    let elem_start = self.chunk.instructions.len();
                     self.name(t);
                     self.infix_bp(0);
+                    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);
+                    }
                 }
             } else {
+                // Support genexpr in arg position (expr for ...)
+                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);
+                }
             }
             argc += 1;
             if matches!(self.peek(), Some(TokenType::Comma)) {

compiler/tests/cases/parser_cases.json

@@ -1646,5 +1646,12 @@
     "names": ["x_1"],
     "instructions": [["LoadConst", 0], ["BuildList", 1], ["StoreName", 0], ["ReturnValue", 0]],
     "annotations": {"x": "tuple"}
+  },
+  {
+    "src": "sum(i for i in range(3))",
+    "constants": ["3"],
+    "names": ["i_0", "i_1"],
+    "instructions": [["BuildList",0], ["LoadConst",0], ["CallRange",1], ["GetIter",0], ["ForIter",9], ["StoreName",1], ["LoadName",1], ["ListAppend",0], ["Jump",4], ["CallSum",1], ["PopTop",0], ["ReturnValue",0]],
+    "annotations": {}
   }
 ]

compiler/tests/cases/vm_cases.json

@@ -232,5 +232,14 @@
   {"src": "pairs = [(1, 'a'), (2, 'b'), (3, 'c')]\nprint({k: v for k, v in pairs})", "output": ["{1: 'a', 2: 'b', 3: 'c'}"], "result": "None"},
   {"src": "print(sum((x for x in range(5))))", "output": ["10"], "result": "None"},
   {"src": "print([x for x in [1, 2, 3, 4, 5] if x > 2])", "output": ["[3, 4, 5]"], "result": "None"},
-  {"src": "matrix = [[1, 2], [3, 4], [5, 6]]\nprint([row[0] for row in matrix])", "output": ["[1, 3, 5]"], "result": "None"}
+  {"src": "matrix = [[1, 2], [3, 4], [5, 6]]\nprint([row[0] for row in matrix])", "output": ["[1, 3, 5]"], "result": "None"},
+  {"src": "x = [1,2,3]\ny = [4,5,6]\nprint(sum(xi*yi for xi,yi in zip(x,y)))", "output": ["32"], "result": "None"},
+  {"src": "print(sum(i for i in range(5)))", "output": ["10"], "result": "None"},
+  {"src": "print(sum(i*i for i in range(4)))", "output": ["14"], "result": "None"},
+  {"src": "print(list(x*2 for x in range(3)))", "output": ["[0, 2, 4]"], "result": "None"},
+  {"src": "print(max(i for i in [3,1,4,1,5]))", "output": ["5"], "result": "None"},
+  {"src": "print(min(i for i in [3,1,4,1,5]))", "output": ["1"], "result": "None"},
+  {"src": "print(sum(x for x in range(10) if x % 2 == 0))", "output": ["20"], "result": "None"},
+  {"src": "print(sum(i*j for i in range(3) for j in range(3)))", "output": ["9"], "result": "None"},
+  {"src": "print(sum(x+y for x,y in zip([1,2],[3,4])))", "output": ["10"], "result": "None"}
 ]