From 4dc3f08fa08a7216676fd2285bfaa4ad01f28797 Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Sun, 24 Jul 2022 15:54:36 +0900
Subject: [PATCH] Fix upvalue capture in generators

---
 src/obj_gen.c                         | 38 +++++++++++++++++++++++
 test/testGeneratorUpvalues.krk        | 43 +++++++++++++++++++++++++++
 test/testGeneratorUpvalues.krk.expect | 13 ++++++++
 3 files changed, 94 insertions(+)
 create mode 100644 test/testGeneratorUpvalues.krk
 create mode 100644 test/testGeneratorUpvalues.krk.expect

diff --git a/src/obj_gen.c b/src/obj_gen.c
index 5857126..078b1fd 100644
--- a/src/obj_gen.c
+++ b/src/obj_gen.c
@@ -28,6 +28,8 @@ struct generator {
 	int started;
 	KrkValue result;
 	int type;
+	KrkThreadState fakethread;
+	KrkUpvalue * capturedUpvalues;
 };
 
 #define AS_generator(o) ((struct generator *)AS_OBJECT(o))
@@ -36,21 +38,35 @@ struct generator {
 #define CURRENT_CTYPE struct generator *
 #define CURRENT_NAME  self
 
+static void _generator_close_upvalues(struct generator * self) {
+	while (self->capturedUpvalues) {
+		KrkUpvalue * upvalue = self->capturedUpvalues;
+		upvalue->closed = self->args[upvalue->location];
+		upvalue->location = -1;
+		self->capturedUpvalues = upvalue->next;
+	}
+}
+
 static void _generator_gcscan(KrkInstance * _self) {
 	struct generator * self = (struct generator*)_self;
 	krk_markObject((KrkObj*)self->closure);
 	for (size_t i = 0; i < self->argCount; ++i) {
 		krk_markValue(self->args[i]);
 	}
+	for (KrkUpvalue * upvalue = self->capturedUpvalues; upvalue; upvalue = upvalue->next) {
+		krk_markObject((KrkObj*)upvalue);
+	}
 	krk_markValue(self->result);
 }
 
 static void _generator_gcsweep(KrkInstance * self) {
+	_generator_close_upvalues((struct generator*)self);
 	free(((struct generator*)self)->args);
 }
 
 static void _set_generator_done(struct generator * self) {
 	self->ip = NULL;
+	_generator_close_upvalues(self);
 }
 
 /**
@@ -120,12 +136,23 @@ KRK_Method(generator,__call__) {
 	frame->slots   = krk_currentThread.stackTop - krk_currentThread.stack;
 	frame->outSlots = frame->slots;
 	frame->globals = self->closure->globalsTable;
+	frame->globalsOwner = self->closure->globalsOwner;
 
 	/* Stick our stack on their stack */
 	for (size_t i = 0; i < self->argCount; ++i) {
 		krk_push(self->args[i]);
 	}
 
+	/* Point any of our captured upvalues back to their actual stack locations */
+	while (self->capturedUpvalues) {
+		KrkUpvalue * upvalue = self->capturedUpvalues;
+		upvalue->owner = &krk_currentThread;
+		upvalue->location = upvalue->location + frame->slots;
+		self->capturedUpvalues = upvalue->next;
+		upvalue->next = krk_currentThread.openUpvalues;
+		krk_currentThread.openUpvalues = upvalue;
+	}
+
 	if (self->started) {
 		krk_pop();
 		if (argc > 1) {
@@ -157,6 +184,16 @@ KRK_Method(generator,__call__) {
 		return NONE_VAL();
 	}
 
+	/* Redirect any remaining upvalues captured from us, and release them from the VM */
+	while (krk_currentThread.openUpvalues != NULL && krk_currentThread.openUpvalues->location >= (int)frame->slots) {
+		KrkUpvalue * upvalue = krk_currentThread.openUpvalues;
+		upvalue->location = upvalue->location - frame->slots;
+		upvalue->owner = &self->fakethread;
+		krk_currentThread.openUpvalues = upvalue->next;
+		upvalue->next = self->capturedUpvalues;
+		self->capturedUpvalues = upvalue;
+	}
+
 	/* Determine the stack state */
 	if (stackAfter > stackBefore) {
 		size_t newArgs = stackAfter - stackBefore;
@@ -171,6 +208,7 @@ KRK_Method(generator,__call__) {
 	/* Save stack entries */
 	memcpy(self->args, krk_currentThread.stackTop - self->argCount, sizeof(KrkValue) * self->argCount);
 	self->ip      = frame->ip;
+	self->fakethread.stack = self->args;
 
 	krk_currentThread.stackTop = krk_currentThread.stack + frame->slots;
 
diff --git a/test/testGeneratorUpvalues.krk b/test/testGeneratorUpvalues.krk
new file mode 100644
index 0000000..b056552
--- /dev/null
+++ b/test/testGeneratorUpvalues.krk
@@ -0,0 +1,43 @@
+let t = []
+
+def foo():
+    let a, b, c = 1,2,3
+    t.append(lambda: (a,b,c))
+    yield
+    a = 'a'
+    t.append(lambda: (b,c,a))
+    yield
+    b = 'b'
+    t.append(lambda: (c,b,a))
+    yield
+    c = 'c'
+
+let g = foo()
+print('one')
+g()
+print(t[0]())
+if True:
+    print('two')
+    g()
+print(t[0]())
+print(t[1]())
+print('three')
+g()
+print(t[0]())
+print(t[1]())
+print(t[2]())
+print("exhaust")
+g()
+
+del g
+
+import gc
+gc.collect()
+gc.collect()
+gc.collect()
+gc.collect()
+
+
+print(t[0]())
+print(t[1]())
+print(t[2]())
diff --git a/test/testGeneratorUpvalues.krk.expect b/test/testGeneratorUpvalues.krk.expect
new file mode 100644
index 0000000..79a9a8d
--- /dev/null
+++ b/test/testGeneratorUpvalues.krk.expect
@@ -0,0 +1,13 @@
+one
+(1, 2, 3)
+two
+('a', 2, 3)
+(2, 3, 'a')
+three
+('a', 'b', 3)
+('b', 3, 'a')
+(3, 'b', 'a')
+exhaust
+('a', 'b', 'c')
+('b', 'c', 'a')
+('c', 'b', 'a')