Support chained 'for's in generator expressions

src/compiler.c

@@ -2252,23 +2252,8 @@ static void singleInner(ssize_t indLoopCounter) {
 	expression();
 }
 
-/**
- * @brief Parser a generator expression.
- *
- * This is essentially the same as a comprehension body, but
- * without the added loop index or wrapping function call.
- *
- * After every inner expression, we yield.
- */
-static void generatorExpression(KrkScanner scannerBefore, Parser parserBefore) {
-	parser.previous = syntheticToken("<genexpr>");
-	Compiler subcompiler;
-	initCompiler(&subcompiler, TYPE_FUNCTION);
-	subcompiler.function->chunk.filename = subcompiler.enclosing->function->chunk.filename;
-	subcompiler.function->isGenerator = 1;
-
-	beginScope();
 
+static void generatorInner(KrkScanner scannerBefore, Parser parserBefore) {
 	ssize_t loopInd = current->localCount;
 	ssize_t varCount = 0;
 	do {
@@ -2320,21 +2305,43 @@ static void generatorExpression(KrkScanner scannerBefore, Parser parserBefore) {
 		emitByte(OP_POP); /* Pop condition */
 	}
 
-	KrkScanner scannerAfter = krk_tellScanner();
-	Parser  parserAfter = parser;
-	krk_rewindScanner(scannerBefore);
-	parser = parserBefore;
-
 	beginScope();
-	expression();
+	if (match(TOKEN_FOR)) {
-	emitByte(OP_YIELD);
+		generatorInner(scannerBefore, parserBefore);
+	} else {
+		KrkScanner scannerAfter = krk_tellScanner();
+		Parser  parserAfter = parser;
+		krk_rewindScanner(scannerBefore);
+		parser = parserBefore;
+		expression();
+		emitByte(OP_YIELD);
+		krk_rewindScanner(scannerAfter);
+		parser = parserAfter;
+	}
 	endScope();
 
-	krk_rewindScanner(scannerAfter);
-	parser = parserAfter;
-
 	emitLoop(loopStart);
 	patchJump(exitJump);
+}
+
+/**
+ * @brief Parser a generator expression.
+ *
+ * This is essentially the same as a comprehension body, but
+ * without the added loop index or wrapping function call.
+ *
+ * After every inner expression, we yield.
+ */
+static void generatorExpression(KrkScanner scannerBefore, Parser parserBefore) {
+	parser.previous = syntheticToken("<genexpr>");
+	Compiler subcompiler;
+	initCompiler(&subcompiler, TYPE_FUNCTION);
+	subcompiler.function->chunk.filename = subcompiler.enclosing->function->chunk.filename;
+	subcompiler.function->isGenerator = 1;
+
+	beginScope();
+	generatorInner(scannerBefore, parserBefore);
+	endScope();
 
 	KrkFunction *subfunction = endCompiler();
 	size_t indFunc = krk_addConstant(currentChunk(), OBJECT_VAL(subfunction));

test/testGeneratorExpressions.krk

@@ -19,3 +19,5 @@ try:
 except SyntaxError as e:
     print('parenthesized' in str(e))
 
+let matrix = [[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]]
+print(list(entry for row in matrix for entry in row))

test/testGeneratorExpressions.krk.expect

@@ -3,3 +3,4 @@
 True
 285
 True
+[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]