From ee1420cd3578aab59f95bfce01716c059f7fc4df Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Tue, 5 Jul 2022 14:45:47 +0900
Subject: [PATCH] Use krk_unpackIterable (and a temporary tuple) in OP_UNPACK

---
 src/vm.c                           | 92 +++++++++++-------------------
 test/testIteratorUnpack.krk.expect |  4 +-
 2 files changed, 35 insertions(+), 61 deletions(-)

diff --git a/src/vm.c b/src/vm.c
index 381db22..e813544 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -609,6 +609,19 @@ static void multipleDefs(const KrkClosure * closure, int destination) {
 				"<unnamed>")));
 }
 
+static int _unpack_op(void * context, const KrkValue * values, size_t count) {
+	KrkTuple * output = context;
+	if (unlikely(output->values.count + count > output->values.capacity)) {
+		krk_runtimeError(vm.exceptions->valueError, "too many values to unpack (expected %zu)",
+			output->values.capacity);
+		return 1;
+	}
+	for (size_t i = 0; i < count; ++i) {
+		output->values.values[output->values.count++] = values[i];
+	}
+	return 0;
+}
+
 static int _unpack_args(void * context, const KrkValue * values, size_t count) {
 	KrkValueArray * positionals = context;
 	if (positionals->count + count > positionals->capacity) {
@@ -3295,66 +3308,27 @@ _finishReturn: (void)0;
 				ONE_BYTE_OPERAND;
 				size_t count = OPERAND;
 				KrkValue sequence = krk_peek(0);
-				/* First figure out what it is and if we can unpack it. */
-#define unpackArray(counter, indexer) do { \
-					if (counter != count) { \
-						krk_runtimeError(vm.exceptions->valueError, "Wrong number of values to unpack (wanted %d, got %d)", (int)count, (int)counter); \
-						break; \
-					} \
-					for (size_t i = 1; i < counter; ++i) { \
-						krk_push(indexer); \
-					} \
-					size_t i = 0; \
-					krk_currentThread.stackTop[-count] = indexer; \
-				} while (0)
-				if (IS_TUPLE(sequence)) {
-					unpackArray(AS_TUPLE(sequence)->values.count, AS_TUPLE(sequence)->values.values[i]);
-				} else if (IS_INSTANCE(sequence) && AS_INSTANCE(sequence)->_class == vm.baseClasses->listClass) {
-					unpackArray(AS_LIST(sequence)->count, AS_LIST(sequence)->values[i]);
-				} else if (IS_INSTANCE(sequence) && AS_INSTANCE(sequence)->_class == vm.baseClasses->dictClass) {
-					unpackArray(AS_DICT(sequence)->count, krk_dict_nth_key_fast(AS_DICT(sequence)->capacity, AS_DICT(sequence)->entries, i));
-				} else if (IS_STRING(sequence)) {
-					unpackArray(AS_STRING(sequence)->codesLength, krk_string_get(2,(KrkValue[]){sequence,INTEGER_VAL(i)},0));
-				} else {
-					KrkClass * type = krk_getType(sequence);
-					if (!type->_iter) {
-						krk_runtimeError(vm.exceptions->typeError, "'%s' object is not iterable", krk_typeName(sequence));
-						goto _finishException;
-					} else {
-						size_t stackStart = krk_currentThread.stackTop - krk_currentThread.stack - 1;
-						size_t counter = 0;
-						for (size_t i = 0; i < count-1; i++) {
-							krk_push(NONE_VAL());
-						}
-						/* Create the iterator */
-						krk_push(krk_currentThread.stack[stackStart]);
-						krk_push(krk_callDirect(type->_iter, 1));
-
-						do {
-							/* Call it until it gives us itself */
-							krk_push(krk_currentThread.stackTop[-1]);
-							krk_push(krk_callStack(0));
-							if (krk_valuesSame(krk_currentThread.stackTop[-2], krk_currentThread.stackTop[-1])) {
-								/* We're done. */
-								krk_pop(); /* The result of iteration */
-								krk_pop(); /* The iterator */
-								if (counter != count) {
-									krk_runtimeError(vm.exceptions->valueError, "Wrong number of values to unpack (wanted %d, got %d)", (int)count, (int)counter);
-									goto _finishException;
-								}
-								break;
-							}
-							if (counter == count) {
-								krk_runtimeError(vm.exceptions->valueError, "Wrong number of values to unpack (wanted %d, got %d)", (int)count, (int)counter);
-								goto _finishException;
-							}
-							/* Rotate */
-							krk_currentThread.stack[stackStart+counter] = krk_pop();
-							counter++;
-						} while (1);
-					}
+				KrkTuple * values = krk_newTuple(count);
+				krk_push(OBJECT_VAL(values));
+				if (unlikely(krk_unpackIterable(sequence, values, _unpack_op))) {
+					goto _finishException;
 				}
-#undef unpackArray
+				if (unlikely(values->values.count != count)) {
+					krk_runtimeError(vm.exceptions->valueError, "not enough values to unpack (expected %zu, got %zu)", count, values->values.count);
+					goto _finishException;
+				}
+				if (unlikely(count == 0)) {
+					krk_pop();
+					krk_pop();
+					break;
+				}
+				/* We no longer need the sequence */
+				krk_swap(1);
+				krk_pop();
+				for (size_t i = 1; i < values->values.count; ++i) {
+					krk_push(values->values.values[i]);
+				}
+				krk_currentThread.stackTop[-count] = values->values.values[0];
 				break;
 			}
 		}
diff --git a/test/testIteratorUnpack.krk.expect b/test/testIteratorUnpack.krk.expect
index e63e9d0..96c6cdf 100644
--- a/test/testIteratorUnpack.krk.expect
+++ b/test/testIteratorUnpack.krk.expect
@@ -1,4 +1,4 @@
 0 1 2
 0 1 2
-ValueError('Wrong number of values to unpack (wanted 2, got 2)')
-ValueError('Wrong number of values to unpack (wanted 4, got 3)')
+ValueError('too many values to unpack (expected 2)')
+ValueError('not enough values to unpack (expected 4, got 3)')