diff --git a/compiler.c b/compiler.c index 4a87b25..5b446cb 100644 --- a/compiler.c +++ b/compiler.c @@ -1610,13 +1610,17 @@ static void list(int canAssign) { /* x in... */ ssize_t loopInd = current->localCount; - varDeclaration(); - defineVariable(loopInd); + ssize_t varCount = 0; + do { + varDeclaration(); + defineVariable(loopInd); + varCount++; + } while (match(TOKEN_COMMA)); consume(TOKEN_IN, "Only iterator loops (for ... in ...) are allowed in list comprehensions."); beginScope(); - expression(); + parsePrecedence(PREC_OR); /* Otherwise we can get trapped on a ternary */ endScope(); /* iterable... */ @@ -1649,10 +1653,29 @@ static void list(int canAssign) { * and there's no feasible way they can return themselves without * our intended sentinel meaning, right? Surely? */ EMIT_CONSTANT_OP(OP_GET_LOCAL, indLoopIter); - emitBytes(OP_EQUAL, OP_NOT); - int exitJump = emitJump(OP_JUMP_IF_FALSE); + emitByte(OP_EQUAL); + int exitJump = emitJump(OP_JUMP_IF_TRUE); emitByte(OP_POP); + /* Unpack tuple */ + if (varCount > 1) { + EMIT_CONSTANT_OP(OP_GET_LOCAL, loopInd); + EMIT_CONSTANT_OP(OP_UNPACK_TUPLE, varCount); + for (ssize_t i = loopInd + varCount - 1; i >= loopInd; i--) { + EMIT_CONSTANT_OP(OP_SET_LOCAL, i); + emitByte(OP_POP); + } + } + + if (match(TOKEN_IF)) { + parsePrecedence(PREC_OR); + int acceptJump = emitJump(OP_JUMP_IF_TRUE); + emitByte(OP_POP); /* Pop condition */ + emitLoop(loopStart); + patchJump(acceptJump); + emitByte(OP_POP); /* Pop condition */ + } + /* Now we can rewind the scanner to have it parse the original * expression that uses our iterated values! */ KrkScanner scannerAfter = krk_tellScanner(); diff --git a/test/testUnpackIter.krk b/test/testUnpackIter.krk index a2c66f9..789669d 100644 --- a/test/testUnpackIter.krk +++ b/test/testUnpackIter.krk @@ -29,3 +29,12 @@ for i in {'a': 1, 2: 'b'}.items(): print(i) # ('a', 1) # (2, 'b') + +print(['{}: {}'.format(repr(k),repr(v)) for k, v in {1: 2, 3: 4}.items()]) +# ['1: 2', '3: 4'] + +print([x for x,y in [(1,2),(3,4),(5,6)] if x + y > 3]) +# [3, 5] + +print(', '.join(['{}: {}'.format(repr(k),repr(v)) for k, v in {1: 2, '3': 4, 5: 'six', 'seven': 8}.items() if isinstance(k,str)])) +# '3': 4, 'seven': 8 diff --git a/test/testUnpackIter.krk.expect b/test/testUnpackIter.krk.expect index 7fd417d..346e997 100644 --- a/test/testUnpackIter.krk.expect +++ b/test/testUnpackIter.krk.expect @@ -9,3 +9,6 @@ a 1 2 b ('a', 1) (2, 'b') +['1: 2', '3: 4'] +[3, 5] +'3': 4, 'seven': 8