...

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);
+void pk_newgenerator(py_Ref out, Frame* frame);

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){
+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() {

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;

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;