diff --git a/chunk.h b/chunk.h index c1ffee8..c21586b 100644 --- a/chunk.h +++ b/chunk.h @@ -71,6 +71,7 @@ typedef enum { OP_EXPAND_ARGS, OP_FINALIZE, OP_TUPLE, + OP_UNPACK_TUPLE, OP_CONSTANT_LONG = 128, OP_DEFINE_GLOBAL_LONG, @@ -91,6 +92,7 @@ typedef enum { OP_INC_LONG, OP_KWARGS_LONG, OP_TUPLE_LONG, + OP_UNPACK_TUPLE_LONG, } KrkOpCode; typedef struct { diff --git a/compiler.c b/compiler.c index c083256..4a87b25 100644 --- a/compiler.c +++ b/compiler.c @@ -1174,7 +1174,11 @@ static void forStatement() { beginScope(); ssize_t loopInd = current->localCount; - varDeclaration(); + ssize_t varCount = 0; + do { + varDeclaration(); + varCount++; + } while (match(TOKEN_COMMA)); int loopStart; int exitJump; @@ -1212,17 +1216,27 @@ static void forStatement() { /* Get the loop iterator again */ EMIT_CONSTANT_OP(OP_GET_LOCAL, indLoopIter); - emitBytes(OP_EQUAL, OP_NOT); - exitJump = emitJump(OP_JUMP_IF_FALSE); + emitByte(OP_EQUAL); + exitJump = emitJump(OP_JUMP_IF_TRUE); emitByte(OP_POP); + 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); + } + } + } else { consume(TOKEN_SEMICOLON,"expect ; after var declaration in for loop"); loopStart = currentChunk()->count; - beginScope(); - expression(); /* condition */ + do { + expression(); /* condition */ + } while (match(TOKEN_COMMA)); endScope(); exitJump = emitJump(OP_JUMP_IF_FALSE); emitByte(OP_POP); @@ -1232,7 +1246,9 @@ static void forStatement() { int bodyJump = emitJump(OP_JUMP); int incrementStart = currentChunk()->count; beginScope(); - expression(); + do { + expression(); + } while (match(TOKEN_COMMA)); endScope(); emitByte(OP_POP); diff --git a/debug.c b/debug.c index 173a679..195f0b9 100644 --- a/debug.c +++ b/debug.c @@ -140,6 +140,7 @@ size_t krk_disassembleInstruction(FILE * f, KrkFunction * func, size_t offset) { OPERAND(OP_CALL, (void)0) OPERAND(OP_INC, (void)0) OPERAND(OP_TUPLE, (void)0) + OPERAND(OP_UNPACK_TUPLE, (void)0) JUMP(OP_JUMP,+) JUMP(OP_JUMP_IF_FALSE,+) JUMP(OP_JUMP_IF_TRUE,+) diff --git a/test/testUnpackIter.krk b/test/testUnpackIter.krk new file mode 100644 index 0000000..a2c66f9 --- /dev/null +++ b/test/testUnpackIter.krk @@ -0,0 +1,31 @@ +for k,v in [(1,2),(3,4),(5,6)]: + print(k,v) +# 1 2 +# 3 4 +# 5 6 + +try: + for k,v,z in [(1,2,7),(3,4,8),(5,6)]: + print(k,v,z) +except: + print(exception.__class__.__name__) +# 1 2 7 +# 3 4 8 +# ValueError + +try: + for k,v in [1,2,3]: + print("type error") +except: + print(exception.__class__.__name__) +# TypeError + +for k, v in {'a': 1, 2: 'b'}.items(): + print(k,v) +# a 1 +# 2 b + +for i in {'a': 1, 2: 'b'}.items(): + print(i) +# ('a', 1) +# (2, 'b') diff --git a/test/testUnpackIter.krk.expect b/test/testUnpackIter.krk.expect new file mode 100644 index 0000000..7fd417d --- /dev/null +++ b/test/testUnpackIter.krk.expect @@ -0,0 +1,11 @@ +1 2 +3 4 +5 6 +1 2 7 +3 4 8 +ValueError +TypeError +a 1 +2 b +('a', 1) +(2, 'b') diff --git a/vm.c b/vm.c index 929a13e..335fb33 100644 --- a/vm.c +++ b/vm.c @@ -3559,6 +3559,24 @@ static KrkValue run() { } break; } + case OP_UNPACK_TUPLE_LONG: + case OP_UNPACK_TUPLE: { + size_t count = readBytes(frame, operandWidth); + KrkValue tuple = krk_peek(0); + if (!IS_TUPLE(tuple)) { + krk_runtimeError(vm.exceptions.typeError, "Can not unpack non-tuple '%s'", krk_typeName(tuple)); + goto _finishException; + } else if (AS_TUPLE(tuple)->values.count != count) { + krk_runtimeError(vm.exceptions.valueError, "Wrong number of values to unpack (wanted %d, got %d)", (int)count, (int)AS_TUPLE(tuple)->values.count); + goto _finishException; + } + /* Unpack from 1 to end, then unpack 0 into bottom slot */ + for (size_t i = 1; i < AS_TUPLE(tuple)->values.count; ++i) { + krk_push(AS_TUPLE(tuple)->values.values[i]); + } + vm.stackTop[-count] = AS_TUPLE(tuple)->values.values[0]; + break; + } } if (!(vm.flags & KRK_HAS_EXCEPTION)) continue; _finishException: