fix goto issues

src/codeobject.h

@@ -116,6 +116,7 @@ class Frame {
 private:
     std::vector<PyVar> s_data;
     int ip = 0;
+    std::stack<int> forLoops;       // record the FOR_ITER bytecode index
 public:
     PyVar _module;
     PyVarDict f_locals;
@@ -152,10 +153,6 @@ public:
         return v;
     }
 
-    void __clearDataStack(){
-        s_data.clear();
-    }
-
     inline PyVar __deref_pointer(VM*, PyVar);
 
     inline PyVar popValue(VM* vm){
@@ -174,10 +171,36 @@ public:
         s_data.push_back(v);
     }
 
-    inline void jumpTo(int i){
+
+    void __reportForIter(){
+        int lastIp = ip - 1;
+        if(forLoops.empty()) forLoops.push(lastIp);
+        else{
+            if(forLoops.top() == lastIp) return;
+            if(forLoops.top() < lastIp) forLoops.push(lastIp);
+            else UNREACHABLE();
+        }
+    }
+
+    inline void jump(int i){
         this->ip = i;
     }
 
+    void safeJump(int i){
+        this->ip = i;
+        while(!forLoops.empty()){
+            int start = forLoops.top();
+            int end = code->co_code[start].arg;
+            if(i < start || i >= end){
+                //printf("%d <- [%d, %d)\n", i, start, end);
+                __pop();    // pop the iterator
+                forLoops.pop();
+            }else{
+                break;
+            }
+        }
+    }
+
     PyVarList popNValuesReversed(VM* vm, int n){
         PyVarList v(n);
         for(int i=n-1; i>=0; i--) v[i] = popValue(vm);

src/compiler.h

@@ -16,10 +16,9 @@ struct GrammarRule{
 };
 
 struct Loop {
-    bool forLoop;
     int start;
     std::vector<int> breaks;
-    Loop(bool forLoop, int start) : forLoop(forLoop), start(start) {}
+    Loop(int start) : start(start) {}
 };
 
 class Compiler {
@@ -484,7 +483,7 @@ __LISTCOMP:
         patchJump(_skipPatch);
 
         emitCode(OP_GET_ITER);
-        Loop& loop = enterLoop(true);
+        Loop& loop = enterLoop();
         int patch = emitCode(OP_FOR_ITER);
 
         if(_cond_end != _cond_start) {      // there is an if condition
@@ -682,8 +681,8 @@ __LISTCOMP:
         }
     }
 
-    Loop& enterLoop(bool forLoop){
+    Loop& enterLoop(){
-        Loop lp(forLoop, (int)getCode()->co_code.size());
+        Loop lp((int)getCode()->co_code.size());
         loops.push(lp);
         return loops.top();
     }
@@ -695,7 +694,7 @@ __LISTCOMP:
     }
 
     void compileWhileLoop() {
-        Loop& loop = enterLoop(false);
+        Loop& loop = enterLoop();
         EXPR_TUPLE();
         int patch = emitCode(OP_POP_JUMP_IF_FALSE);
         compileBlockBody();
@@ -716,7 +715,7 @@ __LISTCOMP:
     void compileForLoop() {
         EXPR_FOR_VARS();consume(TK("in"));EXPR_TUPLE();
         emitCode(OP_GET_ITER);
-        Loop& loop = enterLoop(true);
+        Loop& loop = enterLoop();
         int patch = emitCode(OP_FOR_ITER);
         compileBlockBody();
         emitCode(OP_JUMP_ABSOLUTE, loop.start); keepOpcodeLine();
@@ -728,8 +727,7 @@ __LISTCOMP:
         if (match(TK("break"))) {
             if (loops.empty()) syntaxError("'break' outside loop");
             consumeEndStatement();
-            if(getLoop().forLoop) emitCode(OP_POP_TOP); // pop the iterator of for loop.
+            int patch = emitCode(OP_SAFE_JUMP_ABSOLUTE);
-            int patch = emitCode(OP_JUMP_ABSOLUTE);
             getLoop().breaks.push_back(patch);
         } else if (match(TK("continue"))) {
             if (loops.empty()) syntaxError("'continue' not properly in loop");

src/opcodes.h

@@ -30,6 +30,7 @@ OPCODE(FOR_ITER)
 
 OPCODE(POP_JUMP_IF_FALSE)
 OPCODE(JUMP_ABSOLUTE)
+OPCODE(SAFE_JUMP_ABSOLUTE)
 OPCODE(JUMP_IF_TRUE_OR_POP)
 OPCODE(JUMP_IF_FALSE_OR_POP)
 

src/vm.h

@@ -207,7 +207,7 @@ private:
                     frame->push(PyBool(!PyBool_AS_C(obj_bool)));
                 } break;
             case OP_POP_JUMP_IF_FALSE:
-                if(!PyBool_AS_C(asBool(frame->popValue(this)))) frame->jumpTo(byte.arg);
+                if(!PyBool_AS_C(asBool(frame->popValue(this)))) frame->jump(byte.arg);
                 break;
             case OP_LOAD_NONE: frame->push(None); break;
             case OP_LOAD_TRUE: frame->push(True); break;
@@ -246,7 +246,8 @@ private:
                     if(ret == __py2py_call_signal) return ret;
                     frame->push(ret);
                 } break;
-            case OP_JUMP_ABSOLUTE: frame->jumpTo(byte.arg); break;
+            case OP_JUMP_ABSOLUTE: frame->jump(byte.arg); break;
+            case OP_SAFE_JUMP_ABSOLUTE: frame->safeJump(byte.arg); break;
             case OP_GOTO: {
                 PyVar obj = frame->popValue(this);
                 const _Str& label = PyStr_AS_C(obj);
@@ -254,8 +255,7 @@ private:
                 if(it == frame->code->co_labels.end()){
                     _error("KeyError", "label '" + label + "' not found");
                 }
-                frame->__clearDataStack();
+                frame->safeJump(it->second);
-                frame->jumpTo(it->second);
             } break;
             case OP_GET_ITER:
                 {
@@ -271,26 +271,26 @@ private:
                 } break;
             case OP_FOR_ITER:
                 {
+                    frame->__reportForIter();
                     const PyVar& iter = frame->topValue(this);
                     auto& it = PyIter_AS_C(iter);
                     if(it->hasNext()){
                         it->var->set(this, frame, it->next());
                     }
                     else{
-                        frame->popValue(this);
+                        frame->safeJump(byte.arg);
-                        frame->jumpTo(byte.arg);
                     }
                 } break;
             case OP_JUMP_IF_FALSE_OR_POP:
                 {
                     const PyVar& expr = frame->topValue(this);
-                    if(asBool(expr)==False) frame->jumpTo(byte.arg);
+                    if(asBool(expr)==False) frame->jump(byte.arg);
                     else frame->popValue(this);
                 } break;
             case OP_JUMP_IF_TRUE_OR_POP:
                 {
                     const PyVar& expr = frame->topValue(this);
-                    if(asBool(expr)==True) frame->jumpTo(byte.arg);
+                    if(asBool(expr)==True) frame->jump(byte.arg);
                     else frame->popValue(this);
                 } break;
             case OP_BUILD_SLICE:

tests/5.py

@@ -0,0 +1,9 @@
+for i in range(10):
+    for j in range(10):
+        goto .test 
+        print(2)
+    label .test
+    print(i)
+
+# 15, 23
+# 5, 28