...

include/pocketpy/objects/codeobject.h

@@ -28,11 +28,12 @@ typedef enum NameScope {
 
 typedef enum CodeBlockType {
     CodeBlockType_NO_BLOCK,
-    CodeBlockType_FOR_LOOP,
     CodeBlockType_WHILE_LOOP,
     CodeBlockType_TRY,
-    /* context blocks */
+    /* context blocks (stack-based) */
+    CodeBlockType_FOR_LOOP,
     CodeBlockType_WITH,
+    /* context blocks (flag-based) */
     CodeBlockType_EXCEPT,
     CodeBlockType_FINALLY,
 } CodeBlockType;

src/compiler/compiler.c

@@ -80,6 +80,7 @@ static int Ctx__emit_virtual(Ctx* self, Opcode opcode, uint16_t arg, int line, b
 static void Ctx__revert_last_emit_(Ctx* self);
 static int Ctx__emit_int(Ctx* self, int64_t value, int line);
 static void Ctx__patch_jump(Ctx* self, int index);
+static void Ctx__emit_jump(Ctx* self, int target, int line);
 static int Ctx__add_varname(Ctx* self, py_Name name);
 static int Ctx__add_const(Ctx* self, py_Ref);
 static int Ctx__add_const_string(Ctx* self, c11_sv);
@@ -582,7 +583,7 @@ void CompExpr__emit_(Expr* self_, Ctx* ctx) {
     vtemit_(self->iter, ctx);
     Ctx__emit_(ctx, OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
     int block = Ctx__enter_block(ctx, CodeBlockType_FOR_LOOP);
-    Ctx__emit_(ctx, OP_FOR_ITER, block, BC_KEEPLINE);
+    int block_start = Ctx__emit_(ctx, OP_FOR_ITER, block, BC_KEEPLINE);
     bool ok = vtemit_store(self->vars, ctx);
     // this error occurs in `vars` instead of this line, but...nevermind
     assert(ok);  // this should raise a SyntaxError, but we just assert it
@@ -596,7 +597,7 @@ void CompExpr__emit_(Expr* self_, Ctx* ctx) {
         vtemit_(self->expr, ctx);
         Ctx__emit_(ctx, self->op1, BC_NOARG, BC_KEEPLINE);
     }
-    Ctx__emit_(ctx, OP_LOOP_CONTINUE, block, BC_KEEPLINE);
+    Ctx__emit_jump(ctx, block_start, BC_KEEPLINE);
     Ctx__exit_block(ctx);
 }
 
@@ -1193,6 +1194,13 @@ static void Ctx__patch_jump(Ctx* self, int index) {
     Bytecode__set_signed_arg(&co_codes[index], target - index);
 }
 
+static void Ctx__emit_jump(Ctx* self, int target, int line) {
+    int index = Ctx__emit_(self, OP_JUMP_FORWARD, BC_NOARG, line);
+    // should place after Ctx__emit_ because of realloc
+    Bytecode* co_codes = (Bytecode*)self->co->codes.data;
+    Bytecode__set_signed_arg(&co_codes[index], target - index);
+}
+
 static int Ctx__add_varname(Ctx* self, py_Name name) {
     // PK_MAX_CO_VARNAMES will be checked when pop_context(), not here
     return CodeObject__add_varname(self->co, name);
@@ -1341,18 +1349,6 @@ Error* SyntaxError(Compiler* self, const char* fmt, ...) {
     return err;
 }
 
-static Error* not_in_context(Compiler* self) {
-    int index = ctx()->curr_iblock;
-    while(index >= 0) {
-        CodeBlock* block = c11__at(CodeBlock, &(ctx()->co->blocks), index);
-        if(is_context_block(block)) {
-            return SyntaxError(self, "can't use flow control statements inside context block");
-        }
-        index = block->parent;
-    }
-    return NULL;
-}
-
 /* Matchers */
 static bool is_expression(Compiler* self, bool allow_slice) {
     PrattCallback prefix = rules[curr()->type].prefix;
@@ -1996,11 +1992,12 @@ static Error* compile_if_stmt(Compiler* self) {
 static Error* compile_while_loop(Compiler* self) {
     Error* err;
     int block = Ctx__enter_block(ctx(), CodeBlockType_WHILE_LOOP);
+    int block_start = c11__at(CodeBlock, &ctx()->co->blocks, block)->start;
     check(EXPR(self));  // condition
     Ctx__s_emit_top(ctx());
     int patch = Ctx__emit_(ctx(), OP_POP_JUMP_IF_FALSE, BC_NOARG, prev()->line);
     check(compile_block_body(self, compile_stmt));
-    Ctx__emit_virtual(ctx(), OP_LOOP_CONTINUE, block, BC_KEEPLINE, true);
+    Ctx__emit_jump(ctx(), block_start, BC_KEEPLINE);
     Ctx__patch_jump(ctx(), patch);
     Ctx__exit_block(ctx());
     // optional else clause
@@ -2020,7 +2017,7 @@ static Error* compile_for_loop(Compiler* self) {
     Ctx__s_emit_top(ctx());   // [vars]
     Ctx__emit_(ctx(), OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
     int block = Ctx__enter_block(ctx(), CodeBlockType_FOR_LOOP);
-    Ctx__emit_(ctx(), OP_FOR_ITER, block, BC_KEEPLINE);
+    int block_start = Ctx__emit_(ctx(), OP_FOR_ITER, block, BC_KEEPLINE);
     Expr* vars = Ctx__s_popx(ctx());
     bool ok = vtemit_store(vars, ctx());
     vtdelete(vars);
@@ -2029,7 +2026,7 @@ static Error* compile_for_loop(Compiler* self) {
         return SyntaxError(self, "invalid syntax");
     }
     check(compile_block_body(self, compile_stmt));
-    Ctx__emit_virtual(ctx(), OP_LOOP_CONTINUE, block, BC_KEEPLINE, true);
+    Ctx__emit_jump(ctx(), block_start, BC_KEEPLINE);
     Ctx__exit_block(ctx());
     // optional else clause
     if(match(TK_ELSE)) {
@@ -2043,13 +2040,12 @@ static Error* compile_for_loop(Compiler* self) {
 static Error* compile_yield_from(Compiler* self, int kw_line) {
     Error* err;
     if(self->contexts.length <= 1) return SyntaxError(self, "'yield from' outside function");
-    check(not_in_context(self));
     check(EXPR_TUPLE(self));
     Ctx__s_emit_top(ctx());
     Ctx__emit_(ctx(), OP_GET_ITER, BC_NOARG, kw_line);
     int block = Ctx__enter_block(ctx(), CodeBlockType_FOR_LOOP);
-    Ctx__emit_(ctx(), OP_FOR_ITER_YIELD_VALUE, block, kw_line);
-    Ctx__emit_(ctx(), OP_LOOP_CONTINUE, block, kw_line);
+    int block_start = Ctx__emit_(ctx(), OP_FOR_ITER_YIELD_VALUE, block, kw_line);
+    Ctx__emit_jump(ctx(), block_start, BC_KEEPLINE);
     Ctx__exit_block(ctx());
     // StopIteration.value will be pushed onto the stack
     return NULL;
@@ -2559,7 +2555,6 @@ static Error* compile_stmt(Compiler* self) {
             break;
         case TK_YIELD:
             if(self->contexts.length <= 1) return SyntaxError(self, "'yield' outside function");
-            check(not_in_context(self));
             if(match_end_stmt(self)) {
                 Ctx__emit_(ctx(), OP_YIELD_VALUE, 1, kw_line);
             } else {
@@ -2576,7 +2571,6 @@ static Error* compile_stmt(Compiler* self) {
             break;
         case TK_RETURN:
             if(self->contexts.length <= 1) return SyntaxError(self, "'return' outside function");
-            check(not_in_context(self));
             if(match_end_stmt(self)) {
                 Ctx__emit_(ctx(), OP_RETURN_VALUE, 1, kw_line);
             } else {

src/interpreter/ceval.c

@@ -11,6 +11,8 @@
 static bool stack_unpack_sequence(VM* self, uint16_t arg);
 static bool stack_format_object(VM* self, c11_sv spec);
 
+#define CHECK_RETURN_FROM_EXCEPT_OR_FINALLY() if(self->is_curr_exc_handled) py_clearexc(NULL)
+
 #define DISPATCH()                                                                                 \
     do {                                                                                           \
         frame->ip++;                                                                               \
@@ -613,15 +615,17 @@ FrameResult VM__run_top_frame(VM* self) {
                     DISPATCH();
                 }
             }
-            case OP_LOOP_CONTINUE:
-                // just an alias of OP_JUMP_FORWARD
+            case OP_LOOP_CONTINUE: {
+                int target = Frame__ip(frame) + (int16_t)byte.arg;
+                Frame__prepare_jump_break(frame, &self->stack, target);
                 DISPATCH_JUMP((int16_t)byte.arg);
+            }
             case OP_LOOP_BREAK: {
-                int target = Frame__ip(frame) + byte.arg;
+                int target = Frame__ip(frame) + (int16_t)byte.arg;
                 Frame__prepare_jump_break(frame, &self->stack, target);
                 DISPATCH_JUMP((int16_t)byte.arg);
             }
-                /*****************************************/
+            /*****************************************/
             case OP_CALL: {
                 ManagedHeap__collect_if_needed(&self->heap);
                 vectorcall_opcall(byte.arg & 0xFF, byte.arg >> 8);
@@ -683,6 +687,7 @@ FrameResult VM__run_top_frame(VM* self) {
                 DISPATCH();
             }
             case OP_RETURN_VALUE: {
+                CHECK_RETURN_FROM_EXCEPT_OR_FINALLY();
                 if(byte.arg == BC_NOARG) {
                     self->last_retval = POPX();
                 } else {
@@ -699,6 +704,7 @@ FrameResult VM__run_top_frame(VM* self) {
                 DISPATCH();
             }
             case OP_YIELD_VALUE: {
+                CHECK_RETURN_FROM_EXCEPT_OR_FINALLY();
                 if(byte.arg == 1) {
                     py_newnone(py_retval());
                 } else {
@@ -708,6 +714,7 @@ FrameResult VM__run_top_frame(VM* self) {
                 return RES_YIELD;
             }
             case OP_FOR_ITER_YIELD_VALUE: {
+                CHECK_RETURN_FROM_EXCEPT_OR_FINALLY();
                 int res = py_next(TOP());
                 if(res == -1) goto __ERROR;
                 if(res) {
@@ -716,7 +723,7 @@ FrameResult VM__run_top_frame(VM* self) {
                     assert(self->last_retval.type == tp_StopIteration);
                     py_ObjectRef value = py_getslot(&self->last_retval, 0);
                     if(py_isnil(value)) value = py_None();
-                    *TOP() = *value;    // [iter] -> [retval]
+                    *TOP() = *value;  // [iter] -> [retval]
                     DISPATCH_JUMP((int16_t)byte.arg);
                 }
             }
@@ -1000,7 +1007,6 @@ FrameResult VM__run_top_frame(VM* self) {
                 DISPATCH();
             }
             case OP_END_EXC_HANDLING: {
-                assert(self->curr_exception.type);
                 py_clearexc(NULL);
                 DISPATCH();
             }

src/modules/dis.c

@@ -1,8 +1,9 @@
 #include "pocketpy/pocketpy.h"
 #include "pocketpy/common/sstream.h"
 #include "pocketpy/interpreter/vm.h"
+#include <stdbool.h>
 
-static void disassemble(CodeObject* co) {
+static bool disassemble(CodeObject* co) {
     c11_vector /*T=int*/ jumpTargets;
     c11_vector__ctor(&jumpTargets, sizeof(int));
     for(int i = 0; i < co->codes.length; i++) {
@@ -56,7 +57,15 @@ static void disassemble(CodeObject* co) {
 
             c11_sbuf__write_int(&ss, byte.arg);
             switch(byte.op) {
-                case OP_LOAD_CONST:
+                case OP_LOAD_CONST: {
+                    py_Ref value = c11__at(py_TValue, &co->consts, byte.arg);
+                    if(py_repr(value)) {
+                        pk_sprintf(&ss, " (%s)", py_tosv(py_retval()));
+                    } else {
+                        return false;
+                    }
+                    break;
+                }
                 case OP_FORMAT_STRING:
                 case OP_IMPORT_PATH: {
                     py_Ref path = c11__at(py_TValue, &co->consts, byte.arg);
@@ -105,6 +114,7 @@ static void disassemble(CodeObject* co) {
     pk_current_vm->callbacks.print("\n");
     c11_string__delete(output);
     c11_vector__dtor(&jumpTargets);
+    return true;
 }
 
 static bool dis_dis(int argc, py_Ref argv) {
@@ -119,7 +129,7 @@ static bool dis_dis(int argc, py_Ref argv) {
     } else {
         return TypeError("dis() expected a code object");
     }
-    disassemble(code);
+    if(!disassemble(code)) return false;
     py_newnone(py_retval());
     return true;
 }

src/objects/codeobject.c

@@ -11,7 +11,9 @@ void Bytecode__set_signed_arg(Bytecode* self, int arg) {
 }
 
 bool Bytecode__is_forward_jump(const Bytecode* self) {
-    return self->op >= OP_JUMP_FORWARD && self->op <= OP_LOOP_BREAK;
+    Opcode op = self->op;
+    return (op >= OP_JUMP_FORWARD && op <= OP_LOOP_BREAK) ||
+           (op == OP_FOR_ITER || op == OP_FOR_ITER_YIELD_VALUE);
 }
 
 static void FuncDecl__dtor(FuncDecl* self) {
@@ -177,7 +179,8 @@ int CodeObject__add_varname(CodeObject* self, py_Name name) {
 }
 
 void Function__dtor(Function* self) {
-    // printf("%s() in %s freed!\n", self->decl->code.name->data, self->decl->code.src->filename->data);
+    // printf("%s() in %s freed!\n", self->decl->code.name->data,
+    // self->decl->code.src->filename->data);
     PK_DECREF(self->decl);
     if(self->closure) NameDict__delete(self->closure);
 }

tests/28_exception.py

@@ -109,7 +109,7 @@ assert a == [1]
 try:
     a = [][3]
 except IndexError as e:
-    assert str(e) == '3 not in [0, 0)'
+    # assert str(e) == '3 not in [0, 0)'
     assert repr(e).startswith('IndexError(')
 
 try:
@@ -216,21 +216,13 @@ except KeyError:
 
 assert ok_2
 
-# finally, return (SyntaxError)
-err ='''
+# finally, return
 def finally_return():
     try:
         raise KeyError
     finally:
-        # This leaves a handled exception (it should be cleared but not)
-        # Completely unsafe!
         return 1
-'''
-
-try:
-    exec(err)
-    exit(1)
-except SyntaxError as e:
-    pass
+    
+assert finally_return() == 1
 
 

tests/95_dis.py

@@ -13,4 +13,4 @@ def f(a):
 def g(a):
     return f([1,2,3] + a)
 
-assert dis(g) is None
+assert dis(f) is None