From 070b4bc35695d2ff82490904fca73e8e610edf6b Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Fri, 8 Jan 2021 17:58:33 +0900
Subject: [PATCH] Add tuple unpacking and 'if' conditions to list
 comprehensions

---
 compiler.c                     | 33 ++++++++++++++++++++++++++++-----
 test/testUnpackIter.krk        |  9 +++++++++
 test/testUnpackIter.krk.expect |  3 +++
 3 files changed, 40 insertions(+), 5 deletions(-)

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