From 35b4032de3098ea77e95a2e1d96d785d0e40e895 Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Sat, 6 Mar 2021 15:32:30 +0900
Subject: [PATCH] Support chained 'for's in generator expressions

---
 src/compiler.c                           | 59 +++++++++++++-----------
 test/testGeneratorExpressions.krk        |  2 +
 test/testGeneratorExpressions.krk.expect |  1 +
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/src/compiler.c b/src/compiler.c
index 00af13e..c23ff17 100644
--- a/src/compiler.c
+++ b/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();
-	emitByte(OP_YIELD);
+	if (match(TOKEN_FOR)) {
+		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));
diff --git a/test/testGeneratorExpressions.krk b/test/testGeneratorExpressions.krk
index 3f941ff..3ce1b7e 100644
--- a/test/testGeneratorExpressions.krk
+++ b/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))
diff --git a/test/testGeneratorExpressions.krk.expect b/test/testGeneratorExpressions.krk.expect
index 47d4e97..3dc1f9f 100644
--- a/test/testGeneratorExpressions.krk.expect
+++ b/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]