fix https://github.com/blueloveTH/pocketpy/issues/187

docs/features/ub.md

@@ -7,7 +7,6 @@ These are the undefined behaviours of pkpy. The behaviour of pkpy is undefined i
 
 1. Delete a builtin object. For example, `del int.__add__`.
 2. Call an unbound method with the wrong type of `self`. For example, `int.__add__('1', 2)`.
-3. Use goto statement to jump out of a context block.
-4. Type `T`'s `__new__` returns an object that is not an instance of `T`.
-5. Call `__new__` with a type that is not a subclass of `type`.
-6. `__eq__`, `__lt__` or `__contains__`, etc.. returns a value that is not a boolean.
+3. Type `T`'s `__new__` returns an object that is not an instance of `T`.
+4. Call `__new__` with a type that is not a subclass of `type`.
+5. `__eq__`, `__lt__` or `__contains__`, etc.. returns a value that is not a boolean.

include/pocketpy/codeobject.h

@@ -24,7 +24,7 @@ struct Bytecode{
     uint16_t arg;
 };
 
-enum CodeBlockType {
+enum class CodeBlockType {
     NO_BLOCK,
     FOR_LOOP,
     WHILE_LOOP,
@@ -38,13 +38,13 @@ inline const int BC_KEEPLINE = -1;
 struct CodeBlock {
     CodeBlockType type;
     int parent;         // parent index in blocks
-    int for_loop_depth; // this is used for exception handling
+    int base_stack_size; // this is used for exception handling
     int start;          // start index of this block in codes, inclusive
     int end;            // end index of this block in codes, exclusive
     int end2;           // ...
 
-    CodeBlock(CodeBlockType type, int parent, int for_loop_depth, int start):
-        type(type), parent(parent), for_loop_depth(for_loop_depth), start(start), end(-1), end2(-1) {}
+    CodeBlock(CodeBlockType type, int parent, int base_stack_size, int start):
+        type(type), parent(parent), base_stack_size(base_stack_size), start(start), end(-1), end2(-1) {}
 
     int get_break_end() const{
         if(end2 != -1) return end2;
@@ -68,7 +68,7 @@ struct CodeObject {
     List consts;
     std::vector<StrName> varnames;      // local variables
     NameDictInt varnames_inv;
-    std::vector<CodeBlock> blocks = { CodeBlock(NO_BLOCK, -1, 0, 0) };
+    std::vector<CodeBlock> blocks = { CodeBlock(CodeBlockType::NO_BLOCK, -1, 0, 0) };
     NameDictInt labels;
     std::vector<FuncDecl_> func_decls;
 

include/pocketpy/compiler.h

@@ -107,7 +107,7 @@ class Compiler {
     void exprSubscr();
     void exprLiteral0();
 
-    void compile_block_body();
+    void compile_block_body(void (Compiler::*callback)()=nullptr);
     void compile_normal_import();
     void compile_from_import();
     bool is_expression();

include/pocketpy/expr.h

@@ -25,8 +25,6 @@ struct Expr{
     virtual bool is_name() const { return false; }
     bool is_starred() const { return star_level() > 0; }
 
-    std::string str() const { PK_ASSERT(false); }
-
     // for OP_DELETE_XXX
     [[nodiscard]] virtual bool emit_del(CodeEmitContext* ctx) {
         PK_UNUSED(ctx);
@@ -53,7 +51,7 @@ struct CodeEmitContext{
 
     int curr_block_i = 0;
     bool is_compiling_class = false;
-    int for_loop_depth = 0;
+    int base_stack_size = 0;
 
     std::map<void*, int> _co_consts_nonstring_dedup_map;
     std::map<std::string, int, std::less<>> _co_consts_string_dedup_map;
@@ -62,7 +60,6 @@ struct CodeEmitContext{
     CodeBlock* enter_block(CodeBlockType type);
     void exit_block();
     void emit_expr();   // clear the expression stack and generate bytecode
-    std::string _log_s_expr();
     int emit_(Opcode opcode, uint16_t arg, int line);
     void patch_jump(int index);
     bool add_label(StrName name);

src/ceval.cpp

@@ -782,10 +782,11 @@ __NEXT_STEP:;
     } DISPATCH();
     /*****************************************/
     TARGET(WITH_ENTER)
-        call_method(POPX(), __enter__);
+        PUSH(call_method(TOP(), __enter__));
         DISPATCH();
     TARGET(WITH_EXIT)
-        call_method(POPX(), __exit__);
+        call_method(TOP(), __exit__);
+        POP();
         DISPATCH();
     /*****************************************/
     TARGET(EXCEPTION_MATCH) {

src/compiler.cpp

@@ -25,7 +25,7 @@ namespace pkpy{
 
     void Compiler::pop_context(){
         if(!ctx()->s_expr.empty()){
-            throw std::runtime_error("!ctx()->s_expr.empty()\n" + ctx()->_log_s_expr());
+            throw std::runtime_error("!ctx()->s_expr.empty()");
         }
         // add a `return None` in the end as a guard
         // previously, we only do this if the last opcode is not a return
@@ -491,7 +491,8 @@ __SUBSCR_END:
         ctx()->s_expr.push(make_expr<Literal0Expr>(prev().type));
     }
 
-    void Compiler::compile_block_body() {
+    void Compiler::compile_block_body(void (Compiler::*callback)()) {
+        if(callback == nullptr) callback = &Compiler::compile_stmt;
         consume(TK(":"));
         if(curr().type!=TK("@eol") && curr().type!=TK("@eof")){
             compile_stmt();     // inline block
@@ -503,7 +504,7 @@ __SUBSCR_END:
         consume(TK("@indent"));
         while (curr().type != TK("@dedent")) {
             match_newlines();
-            compile_stmt();
+            (this->*callback)();
             match_newlines();
         }
         consume(TK("@dedent"));
@@ -633,7 +634,7 @@ __EAT_DOTS_END:
     }
 
     void Compiler::compile_while_loop() {
-        CodeBlock* block = ctx()->enter_block(WHILE_LOOP);
+        CodeBlock* block = ctx()->enter_block(CodeBlockType::WHILE_LOOP);
         EXPR(false);   // condition
         int patch = ctx()->emit_(OP_POP_JUMP_IF_FALSE, BC_NOARG, prev().line);
         compile_block_body();
@@ -652,7 +653,7 @@ __EAT_DOTS_END:
         consume(TK("in"));
         EXPR_TUPLE(false);
         ctx()->emit_(OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
-        CodeBlock* block = ctx()->enter_block(FOR_LOOP);
+        CodeBlock* block = ctx()->enter_block(CodeBlockType::FOR_LOOP);
         ctx()->emit_(OP_FOR_ITER, BC_NOARG, BC_KEEPLINE);
         bool ok = vars->emit_store(ctx());
         if(!ok) SyntaxError();  // this error occurs in `vars` instead of this line, but...nevermind
@@ -667,7 +668,7 @@ __EAT_DOTS_END:
     }
 
     void Compiler::compile_try_except() {
-        ctx()->enter_block(TRY_EXCEPT);
+        ctx()->enter_block(CodeBlockType::TRY_EXCEPT);
         compile_block_body();
         std::vector<int> patches = {
             ctx()->emit_(OP_JUMP_ABSOLUTE, BC_NOARG, BC_KEEPLINE)
@@ -813,7 +814,7 @@ __EAT_DOTS_END:
                 // if yield from present, mark the function as generator
                 ctx()->co->is_generator = true;
                 ctx()->emit_(OP_GET_ITER, BC_NOARG, kw_line);
-                ctx()->enter_block(FOR_LOOP);
+                ctx()->enter_block(CodeBlockType::FOR_LOOP);
                 ctx()->emit_(OP_FOR_ITER, BC_NOARG, BC_KEEPLINE);
                 ctx()->emit_(OP_YIELD_VALUE, BC_NOARG, BC_KEEPLINE);
                 ctx()->emit_(OP_LOOP_CONTINUE, ctx()->get_loop(), BC_KEEPLINE);
@@ -909,17 +910,24 @@ __EAT_DOTS_END:
                 consume_end_stmt();
             } break;
             case TK("with"): {
-                EXPR(false);
-                consume(TK("as"));
-                consume(TK("@id"));
-                Expr_ e = make_expr<NameExpr>(prev().str(), name_scope());
-                bool ok = e->emit_store(ctx());
-                if(!ok) SyntaxError();
-                e->emit_(ctx());
+                EXPR(false);    // [ <expr> ]
+                ctx()->enter_block(CodeBlockType::CONTEXT_MANAGER);
+                Expr_ as_name;
+                if(match(TK("as"))){
+                    consume(TK("@id"));
+                    as_name = make_expr<NameExpr>(prev().str(), name_scope());
+                }
                 ctx()->emit_(OP_WITH_ENTER, BC_NOARG, prev().line);
+                // [ <expr> <expr>.__enter__() ]
+                if(as_name){
+                    bool ok = as_name->emit_store(ctx());
+                    if(!ok) SyntaxError();
+                }else{
+                    ctx()->emit_(OP_POP_TOP, BC_NOARG, BC_KEEPLINE);
+                }
                 compile_block_body();
-                e->emit_(ctx());
                 ctx()->emit_(OP_WITH_EXIT, BC_NOARG, prev().line);
+                ctx()->exit_block();
             } break;
             /*************************************************/
             case TK("=="): {

src/expr.cpp

@@ -16,17 +16,17 @@ namespace pkpy{
     int CodeEmitContext::get_loop() const {
         int index = curr_block_i;
         while(index >= 0){
-            if(co->blocks[index].type == FOR_LOOP) break;
-            if(co->blocks[index].type == WHILE_LOOP) break;
+            if(co->blocks[index].type == CodeBlockType::FOR_LOOP) break;
+            if(co->blocks[index].type == CodeBlockType::WHILE_LOOP) break;
             index = co->blocks[index].parent;
         }
         return index;
     }
 
     CodeBlock* CodeEmitContext::enter_block(CodeBlockType type){
-        if(type == FOR_LOOP) for_loop_depth++;
+        if(type==CodeBlockType::FOR_LOOP || type==CodeBlockType::CONTEXT_MANAGER) base_stack_size++;
         co->blocks.push_back(CodeBlock(
-            type, curr_block_i, for_loop_depth, (int)co->codes.size()
+            type, curr_block_i, base_stack_size, (int)co->codes.size()
         ));
         curr_block_i = co->blocks.size()-1;
         return &co->blocks[curr_block_i];
@@ -34,12 +34,12 @@ namespace pkpy{
 
     void CodeEmitContext::exit_block(){
         auto curr_type = co->blocks[curr_block_i].type;
-        if(curr_type == FOR_LOOP) for_loop_depth--;
+        if(curr_type == CodeBlockType::FOR_LOOP || curr_type==CodeBlockType::CONTEXT_MANAGER) base_stack_size--;
         co->blocks[curr_block_i].end = co->codes.size();
         curr_block_i = co->blocks[curr_block_i].parent;
         if(curr_block_i < 0) PK_FATAL_ERROR();
 
-        if(curr_type == FOR_LOOP){
+        if(curr_type == CodeBlockType::FOR_LOOP){
             // add a no op here to make block check work
             emit_(OP_NO_OP, BC_NOARG, BC_KEEPLINE);
         }
@@ -47,19 +47,11 @@ namespace pkpy{
 
     // clear the expression stack and generate bytecode
     void CodeEmitContext::emit_expr(){
-        if(s_expr.size() != 1){
-            throw std::runtime_error("s_expr.size() != 1\n" + _log_s_expr());
-        }
+        if(s_expr.size() != 1) throw std::runtime_error("s_expr.size() != 1");
         Expr_ expr = s_expr.popx();
         expr->emit_(this);
     }
 
-    std::string CodeEmitContext::_log_s_expr(){
-        std::stringstream ss; // debug
-        for(auto& e: s_expr.data()) ss << e->str() << " ";
-        return ss.str();
-    }
-
     int CodeEmitContext::emit_(Opcode opcode, uint16_t arg, int line) {
         co->codes.push_back(Bytecode{(uint8_t)opcode, arg});
         co->iblocks.push_back(curr_block_i);
@@ -379,7 +371,7 @@ namespace pkpy{
         ctx->emit_(op0(), 0, line);
         iter->emit_(ctx);
         ctx->emit_(OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
-        ctx->enter_block(FOR_LOOP);
+        ctx->enter_block(CodeBlockType::FOR_LOOP);
         ctx->emit_(OP_FOR_ITER, BC_NOARG, BC_KEEPLINE);
         bool ok = vars->emit_store(ctx);
         // this error occurs in `vars` instead of this line, but...nevermind

src/frame.cpp

@@ -27,14 +27,14 @@ namespace pkpy{
         // try to find a parent try block
         int block = co->iblocks[_ip];
         while(block >= 0){
-            if(co->blocks[block].type == TRY_EXCEPT) break;
+            if(co->blocks[block].type == CodeBlockType::TRY_EXCEPT) break;
             block = co->blocks[block].parent;
         }
         if(block < 0) return false;
         PyObject* obj = _s->popx();         // pop exception object
-        // get the stack size of the try block (depth of for loops)
-        int _stack_size = co->blocks[block].for_loop_depth;
-        if(stack_size() < _stack_size) throw std::runtime_error("invalid stack size");
+        // get the stack size of the try block
+        int _stack_size = co->blocks[block].base_stack_size;
+        if(stack_size() < _stack_size) throw std::runtime_error(fmt("invalid state: ", stack_size(), '<', _stack_size).str());
         _s->reset(actual_sp_base() + _locals.size() + _stack_size);          // rollback the stack   
         _s->push(obj);                                      // push exception object
         _next_ip = co->blocks[block].end;
@@ -42,7 +42,8 @@ namespace pkpy{
     }
 
     int Frame::_exit_block(int i){
-        if(co->blocks[i].type == FOR_LOOP) _s->pop();
+        auto type = co->blocks[i].type;
+        if(type==CodeBlockType::FOR_LOOP || type==CodeBlockType::CONTEXT_MANAGER) _s->pop();
         return co->blocks[i].parent;
     }
 

src/io.cpp

@@ -91,7 +91,7 @@ unsigned char* _default_import_handler(const char* name_p, int name_size, int* o
             return vm->None;
         });
 
-        vm->bind_method<0>(type, "__enter__", PK_LAMBDA(vm->None));
+        vm->bind_method<0>(type, "__enter__", PK_LAMBDA(args[0]));
     }
 
     FileIO::FileIO(VM* vm, std::string file, std::string mode): file(file), mode(mode) {

tests/27_goto.py

@@ -30,4 +30,9 @@ i += 1
 if i <= 100:
     -> loop
 
-assert sum == 5050
+assert sum == 5050
+
+for i in range(4):
+    _ = 0
+# if there is no op here, the block check will fail
+while i: --i

tests/33_match_case.py

@@ -0,0 +1,44 @@
+# match = 2
+# assert match == 2
+# case = 3
+# assert case == 3
+
+# def f(match):
+#     match match:
+#         case 1: return 1
+#         case 2: return 2
+#         case _:
+#             return 999
+#     return 0
+
+# assert f(1) == 1
+# assert f(2) == 2
+# assert f(3) == 999
+# assert f(4) == 999
+
+# def f():
+#     a = []
+#     try:
+#         match case:
+#             case a[1]: return 1
+#     except IndexError:
+#         return 'IndexError'
+#     return 0
+
+# assert f() == 'IndexError'
+
+
+# def f(pos):
+#     match pos:
+#         case 'str': return 'str'
+#         case 0: return 0
+#         case (1, 2): return '1, 2'
+#         case (3, 4): return '3, 4'
+#         case _: return 'other'
+
+# assert f('str') == 'str'
+# assert f(0) == 0
+# assert f((1, 2)) == '1, 2'
+# assert f((3, 4)) == '3, 4'
+# assert f((1, 3)) == 'other'
+# assert f((1, 2, 3)) == 'other'

tests/34_context.py

@@ -0,0 +1,37 @@
+path = []
+
+class A:
+    def __init__(self, x):
+        self.x = x
+        self.path = []
+
+    def __enter__(self):
+        path.append('enter')
+        return self.x
+    
+    def __exit__(self, *args):
+        path.append('exit')
+    
+
+with A(123):
+    assert path == ['enter']
+assert path == ['enter', 'exit']
+
+path.clear()
+
+with A(123) as a:
+    assert path == ['enter']
+    assert a == 123
+    path.append('in')
+assert path == ['enter', 'in', 'exit']
+
+path.clear()
+
+with A(123) as a:
+    assert path == ['enter']
+    -> end
+    path.append('in')
+
+== end ==
+assert path == ['enter']
+