From f481337f9867499362a7ecf2e1492dfabf1b3273 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Tue, 6 Aug 2024 19:37:44 +0800
Subject: [PATCH] ...

---
 include/pocketpy/interpreter/generator.h |  2 +-
 src/interpreter/generator.c              | 55 ++++++++++++++++++++++--
 src/interpreter/vm.c                     |  4 +-
 src/public/py_exception.c                |  2 -
 tests/{52_yield.py => 51_yield.py}       |  0
 tests/{51_eval.py => 66_eval.py}         |  0
 6 files changed, 54 insertions(+), 9 deletions(-)
 rename tests/{52_yield.py => 51_yield.py} (100%)
 rename tests/{51_eval.py => 66_eval.py} (100%)

diff --git a/include/pocketpy/interpreter/generator.h b/include/pocketpy/interpreter/generator.h
index 6a47d871..e37f623d 100644
--- a/include/pocketpy/interpreter/generator.h
+++ b/include/pocketpy/interpreter/generator.h
@@ -7,4 +7,4 @@ typedef struct Generator{
     int state;
 } Generator;
 
-void pk_newgenerator(py_Ref out, Frame* frame, int slots);
\ No newline at end of file
+void pk_newgenerator(py_Ref out, Frame* frame);
\ No newline at end of file
diff --git a/src/interpreter/generator.c b/src/interpreter/generator.c
index 90993721..f069ed74 100644
--- a/src/interpreter/generator.c
+++ b/src/interpreter/generator.c
@@ -1,16 +1,63 @@
 #include "pocketpy/interpreter/generator.h"
 #include "pocketpy/interpreter/frame.h"
 #include "pocketpy/interpreter/vm.h"
+#include "pocketpy/objects/base.h"
 #include "pocketpy/pocketpy.h"
+#include <stdbool.h>
 
-void pk_newgenerator(py_Ref out, Frame* frame, int slots) {
-    Generator* ud = py_newobject(out, tp_generator, slots, sizeof(Generator));
+void pk_newgenerator(py_Ref out, Frame* frame) {
+    Generator* ud = py_newobject(out, tp_generator, 1, sizeof(Generator));
     ud->frame = frame;
     ud->state = 0;
+    py_newlist(py_getslot(out, 0));
 }
 
-static bool generator__next__(int argc, py_Ref argv){
-    return true;
+static bool generator__next__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    Generator* ud = py_touserdata(argv);
+    py_StackRef p0 = py_peek(0);
+    VM* vm = pk_current_vm;
+    if(ud->state == 2) return StopIteration();
+
+    // reset frame->p0
+    ud->frame->p0 = py_peek(0);
+    ud->frame->locals = py_peek(0);
+
+    // restore the context
+    py_Ref backup = py_getslot(argv, 0);
+    int length = py_list_len(backup);
+    py_TValue* p = py_list_data(backup);
+    for(int i = 0; i < length; i++)
+        py_push(&p[i]);
+    py_list_clear(backup);
+
+    // push frame
+    VM__push_frame(vm, ud->frame);
+
+    FrameResult res = VM__run_top_frame(vm);
+
+    if(res == RES_ERROR) {
+        ud->state = 2;  // end this generator immediately on error
+        if(py_matchexc(tp_StopIteration)) {
+            py_clearexc(p0);
+            return true;
+        }
+        return false;
+    }
+
+    if(res == RES_YIELD) {
+        // backup the context
+        ud->frame = vm->top_frame;
+        for(py_StackRef p = ud->frame->p0; p != vm->stack.sp; p++) {
+            py_list_append(backup, p);
+        }
+        vm->top_frame = vm->top_frame->f_back;
+        ud->state = 1;
+        return true;
+    } else {
+        ud->state = 2;
+        return StopIteration();
+    }
 }
 
 py_Type pk_generator__register() {
diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c
index 4fefe4a9..d25bd95c 100644
--- a/src/interpreter/vm.c
+++ b/src/interpreter/vm.c
@@ -457,7 +457,7 @@ FrameResult VM__vectorcall(VM* self, uint16_t argc, uint16_t kwargc, bool opcall
                 bool ok = prepare_py_call(self->__vectorcall_buffer, argv, p1, kwargc, fn->decl);
                 if(!ok) return RES_ERROR;
                 Frame* frame = Frame__new(co, &fn->module, p0, p0, argv);
-                pk_newgenerator(py_retval(), frame, 0);
+                pk_newgenerator(py_retval(), frame);
                 self->stack.sp = p0;
                 return RES_RETURN;
             }
@@ -592,7 +592,7 @@ void ManagedHeap__mark(ManagedHeap* self) {
 }
 
 void pk_print_stack(VM* self, Frame* frame, Bytecode byte) {
-    return;
+    // return;
     if(frame == NULL || py_isnil(&self->main)) return;
 
     py_TValue* sp = self->stack.sp;
diff --git a/src/public/py_exception.c b/src/public/py_exception.c
index 9ff250a4..6bd1198b 100644
--- a/src/public/py_exception.c
+++ b/src/public/py_exception.c
@@ -54,8 +54,6 @@ static bool _py_BaseException__new__(int argc, py_Ref argv) {
     py_Type cls = py_totype(argv);
     BaseException* ud = py_newobject(py_retval(), cls, 2, sizeof(BaseException));
     c11_vector__ctor(&ud->stacktrace, sizeof(BaseExceptionFrame));
-    py_setslot(py_retval(), 0, py_NIL);
-    py_setslot(py_retval(), 1, py_NIL);
     ud->lineno_backup = -1;
     ud->code_backup = NULL;
     return true;
diff --git a/tests/52_yield.py b/tests/51_yield.py
similarity index 100%
rename from tests/52_yield.py
rename to tests/51_yield.py
diff --git a/tests/51_eval.py b/tests/66_eval.py
similarity index 100%
rename from tests/51_eval.py
rename to tests/66_eval.py