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(""); - 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(""); + 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]