up

src/codeobject.h

@@ -27,6 +27,44 @@ _Str pad(const _Str& s, const int n){
     return s + std::string(n - s.size(), ' ');
 }
 
+enum CodeBlockType {
+    NO_BLOCK,
+    FOR_LOOP,
+    CONTEXT_MANAGER,
+    TRY_EXCEPT,
+};
+
+struct CodeBlock {
+    CodeBlockType type;
+    std::vector<int> id;
+    int parent;        // parent index in co_blocks
+
+    std::string toString() const {
+        if(parent == -1) return "";
+        std::string s = "[";
+        for(int i = 0; i < id.size(); i++){
+            s += std::to_string(id[i]);
+            if(i != id.size()-1) s += "-";
+        }
+        s += ": ";
+        s += std::to_string(type);
+        s += "]";
+        return s;
+    }
+
+    bool operator==(const std::vector<int>& other) const {
+        return id == other;
+    }
+
+    bool operator!=(const std::vector<int>& other) const {
+        return id != other;
+    }
+
+    int depth() const {
+        return id.size();
+    }
+};
+
 struct CodeObject {
     _Source src;
     _Str name;
@@ -45,48 +83,32 @@ struct CodeObject {
     std::vector<std::pair<_Str, NameScope>> co_names;
     std::vector<_Str> co_global_names;
 
-    std::vector<std::vector<int>> co_loops = {{}};
-    int _currLoopIndex = 0;
-
-    std::string getBlockStr(int block){
-        std::vector<int> loopId = co_loops[block];
-        std::string s = "";
-        for(int i=0; i<loopId.size(); i++){
-            s += std::to_string(loopId[i]);
-            if(i != loopId.size()-1) s += "-";
+    std::vector<CodeBlock> co_blocks = { CodeBlock{NO_BLOCK, {}, -1} };
+
+    // tmp variables
+    int _currBlockIndex = 0;
+
+    void __enterBlock(CodeBlockType type){
+        const CodeBlock& currBlock = co_blocks[_currBlockIndex];
+        std::vector<int> copy(currBlock.id);
+        copy.push_back(-1);
+        int t = 0;
+        while(true){
+            copy[copy.size()-1] = t;
+            auto it = std::find(co_blocks.begin(), co_blocks.end(), copy);
+            if(it == co_blocks.end()) break;
+            t++;
         }
-        return s;
+        co_blocks.push_back(CodeBlock{type, copy, _currBlockIndex});
+        _currBlockIndex = co_blocks.size()-1;
     }
 
-    void __enterLoop(int depth){
-        const std::vector<int>& prevLoopId = co_loops[_currLoopIndex];
-        if(depth - prevLoopId.size() == 1){
-            std::vector<int> copy = prevLoopId;
-            copy.push_back(0);
-            int t = 0;
-            while(true){
-                copy[copy.size()-1] = t;
-                auto it = std::find(co_loops.begin(), co_loops.end(), copy);
-                if(it == co_loops.end()) break;
-                t++;
-            }
-            co_loops.push_back(copy);
-        }else{
-            UNREACHABLE();
-        }
-        _currLoopIndex = co_loops.size()-1;
-    }
-
-    void __exitLoop(){
-        std::vector<int> copy = co_loops[_currLoopIndex];
-        copy.pop_back();
-        auto it = std::find(co_loops.begin(), co_loops.end(), copy);
-        if(it == co_loops.end()) UNREACHABLE();
-        _currLoopIndex = it - co_loops.begin();
+    void __exitBlock(){
+        _currBlockIndex = co_blocks[_currBlockIndex].parent;
+        if(_currBlockIndex < 0) UNREACHABLE();
     }
 
     // for goto use
-    // note: some opcodes moves the bytecode, such as listcomp
     // goto/label should be put at toplevel statements
     emhash8::HashMap<_Str, int> co_labels;
 
@@ -125,7 +147,7 @@ public:
     PyVar _module;
     PyVarDict f_locals;
 
-    inline PyVarDict copy_f_locals(){
+    inline PyVarDict copy_f_locals() const {
         return f_locals;
     }
 
@@ -194,27 +216,25 @@ public:
         this->ip += i;
     }
 
-    void jumpAbsoluteSafe(int i){
+    void jumpAbsoluteSafe(int target){
         const ByteCode& prev = code->co_code[this->ip];
-        const std::vector<int> prevLoopId = code->co_loops[prev.block];
-        this->ip = i;
+        int i = prev.block;
+        this->ip = target;
         if(isCodeEnd()){
-            for(int i=0; i<prevLoopId.size(); i++) __pop();
-            return;
-        }
-        const ByteCode& next = code->co_code[i];
-        const std::vector<int> nextLoopId = code->co_loops[next.block];
-        int sizeDelta = prevLoopId.size() - nextLoopId.size();
-        if(sizeDelta < 0){
-            throw std::runtime_error("invalid jump from " + code->getBlockStr(prev.block) + " to " + code->getBlockStr(next.block));
+            while(i>=0){
+                if(code->co_blocks[i].type == FOR_LOOP) __pop();
+                i = code->co_blocks[i].parent;
+            }
         }else{
-            for(int i=0; i<nextLoopId.size(); i++){
-                if(nextLoopId[i] != prevLoopId[i]){
-                    throw std::runtime_error("invalid jump from " + code->getBlockStr(prev.block) + " to " + code->getBlockStr(next.block));
-                }
+            const ByteCode& next = code->co_code[target];
+            while(i>=0 && i!=next.block){
+                if(code->co_blocks[i].type == FOR_LOOP) __pop();
+                i = code->co_blocks[i].parent;
             }
+            if(i!=next.block) throw std::runtime_error(
+                "invalid jump from " + code->co_blocks[prev.block].toString() + " to " + code->co_blocks[next.block].toString()
+            );
         }
-        for(int i=0; i<sizeDelta; i++) __pop();
     }
 
     pkpy::ArgList popNValuesReversed(VM* vm, int n){

src/compiler.h

@@ -575,7 +575,7 @@ __LISTCOMP:
         patchJump(_skipPatch);
 
         emitCode(OP_GET_ITER);
-        Loop& loop = enterLoop();
+        Loop& loop = enterLoop(); getCode()->__enterBlock(FOR_LOOP);
         int patch = emitCode(OP_FOR_ITER);
 
         if(_cond_end_return != -1) {      // there is an if condition
@@ -594,7 +594,7 @@ __LISTCOMP:
 
         emitCode(OP_JUMP_ABSOLUTE, loop.start); keepOpcodeLine();
         patchJump(patch);
-        exitLoop();
+        exitLoop(); getCode()->__exitBlock();
         matchNewLines(mode()==SINGLE_MODE);
         consume(TK("]"));
     }
@@ -710,7 +710,7 @@ __LISTCOMP:
     int emitCode(Opcode opcode, int arg=-1) {
         int line = parser->previous.line;
         getCode()->co_code.push_back(
-            ByteCode{(uint8_t)opcode, arg, (uint16_t)line, (uint16_t)getCode()->_currLoopIndex}
+            ByteCode{(uint8_t)opcode, arg, (uint16_t)line, (uint16_t)getCode()->_currBlockIndex}
         );
         return getCode()->co_code.size() - 1;
     }
@@ -818,14 +818,12 @@ __LISTCOMP:
     }
 
     Loop& enterLoop(){
-        getCode()->__enterLoop(loops.size()+1);
         Loop lp((int)getCode()->co_code.size());
         loops.push(lp);
         return loops.top();
     }
 
     void exitLoop(){
-        getCode()->__exitLoop();
         Loop& lp = loops.top();
         for(int addr : lp.breaks) patchJump(addr);
         loops.pop();
@@ -851,14 +849,14 @@ __LISTCOMP:
     }
 
     void compileForLoop() {
-        EXPR_FOR_VARS();consume(TK("in"));EXPR_TUPLE();
+        EXPR_FOR_VARS();consume(TK("in")); EXPR_TUPLE();
         emitCode(OP_GET_ITER);
-        Loop& loop = enterLoop();
+        Loop& loop = enterLoop(); getCode()->__enterBlock(FOR_LOOP);
         int patch = emitCode(OP_FOR_ITER);
         compileBlockBody();
         emitCode(OP_JUMP_ABSOLUTE, loop.start); keepOpcodeLine();
         patchJump(patch);
-        exitLoop();
+        exitLoop(); getCode()->__exitBlock();
     }
 
     void compileStatement() {

src/opcodes.h

@@ -1,8 +1,15 @@
 #ifdef OPCODE
 
+// Do nothing
 OPCODE(NO_OP)
+
+// This op is a placeholder that should never be executed
 OPCODE(DELETED_OP)
+
+// Load a constant from the `co_consts`
+// ARG: array index
 OPCODE(LOAD_CONST)
+
 OPCODE(IMPORT_NAME)
 OPCODE(PRINT_EXPR)
 OPCODE(POP_TOP)
@@ -67,4 +74,6 @@ OPCODE(WITH_EXIT)
 
 OPCODE(JUMP_RELATIVE)
 
+OPCODE(RAISE_VARARGS)
+
 #endif

src/vm.h

@@ -136,7 +136,6 @@ protected:
                         setAttr(fn, __module__, frame->_module);
                         setAttr(cls, f->name, fn);
                     }
-                    // frame->f_globals()[clsName] = cls;
                 } break;
             case OP_RETURN_VALUE: return frame->popValue(this);
             case OP_PRINT_EXPR:
@@ -548,9 +547,9 @@ public:
         try {
             _Code code = compile(source, filename, mode);
 
-            // if(filename != "<builtins>"){
-            //     std::cout << disassemble(code) << std::endl;
-            // }
+            if(filename != "<builtins>"){
+                std::cout << disassemble(code) << std::endl;
+            }
             
             return _exec(code, _module, {});
         }catch (const _Error& e){
@@ -774,7 +773,6 @@ public:
         int prev_line = -1;
         for(int i=0; i<code->co_code.size(); i++){
             const ByteCode& byte = code->co_code[i];
-            //if(byte.op == OP_NO_OP || byte.op == OP_DELETED_OP) continue;
             _Str line = std::to_string(byte.line);
             if(byte.line == prev_line) line = "";
             else{
@@ -784,7 +782,7 @@ public:
             ss << pad(line, 12) << " " << pad(std::to_string(i), 3);
             ss << " " << pad(OP_NAMES[byte.op], 20) << " ";
             ss << pad(byte.arg == -1 ? "" : std::to_string(byte.arg), 5);
-            ss << '[' << code->getBlockStr(byte.block) << ']';
+            ss << code->co_blocks[byte.block].toString();
             if(i != code->co_code.size() - 1) ss << '\n';
         }
         _StrStream consts;