...

include/pocketpy/interpreter/frame.h

@@ -57,7 +57,6 @@ py_TValue* Frame__f_closure_try_get(Frame* self, py_Name name);
 
 int Frame__prepare_jump_exception_handler(Frame* self, ValueStack*);
 void Frame__prepare_jump_break(Frame* self, ValueStack*, int);
-int Frame__prepare_loop_break(Frame* self, ValueStack*);
 int Frame__exit_block(Frame* self, ValueStack*, int);
 
 UnwindTarget* Frame__find_unwind_target(Frame* self, int iblock);

include/pocketpy/interpreter/vm.h

@@ -59,6 +59,8 @@ py_TypeInfo* pk__type_info(py_Type type);
 bool pk_wrapper__self(int argc, py_Ref argv);
 bool pk_wrapper__NotImplementedError(int argc, py_Ref argv);
 
+const char* pk_op2str(py_Name op);
+
 typedef enum FrameResult {
     RES_ERROR = 0,
     RES_RETURN = 1,

include/pocketpy/objects/codeobject.h

@@ -30,8 +30,11 @@ typedef enum CodeBlockType {
     CodeBlockType_NO_BLOCK,
     CodeBlockType_FOR_LOOP,
     CodeBlockType_WHILE_LOOP,
-    CodeBlockType_CONTEXT_MANAGER,
-    CodeBlockType_TRY_EXCEPT,
+    CodeBlockType_TRY,
+    /* context blocks */
+    CodeBlockType_WITH,
+    CodeBlockType_EXCEPT,
+    CodeBlockType_FINALLY,
 } CodeBlockType;
 
 typedef enum Opcode {

include/pocketpy/xmacros/opcodes.h

@@ -66,8 +66,10 @@ OPCODE(LOOP_BREAK)
 /**************************/
 OPCODE(CALL)
 OPCODE(CALL_VARGS)
+/**************************/
 OPCODE(RETURN_VALUE)
 OPCODE(YIELD_VALUE)
+OPCODE(FOR_ITER_YIELD_VALUE)
 /**************************/
 OPCODE(LIST_APPEND)
 OPCODE(DICT_ADD)
@@ -80,7 +82,6 @@ OPCODE(UNARY_INVERT)
 /**************************/
 OPCODE(GET_ITER)
 OPCODE(FOR_ITER)
-OPCODE(FOR_ITER_YIELD_VALUE)
 /**************************/
 OPCODE(IMPORT_PATH)
 OPCODE(POP_IMPORT_STAR)
@@ -104,6 +105,8 @@ OPCODE(RE_RAISE)
 OPCODE(PUSH_EXCEPTION)
 OPCODE(BEGIN_EXC_HANDLING)
 OPCODE(END_EXC_HANDLING)
+OPCODE(BEGIN_FINALLY)
+OPCODE(END_FINALLY)
 /**************************/
 OPCODE(FORMAT_STRING)
 /**************************/

src/compiler/compiler.c

@@ -72,8 +72,8 @@ typedef struct Expr Expr;
 
 static void Ctx__ctor(Ctx* self, CodeObject* co, FuncDecl* func, int level);
 static void Ctx__dtor(Ctx* self);
-static int Ctx__get_loop(Ctx* self);
-static CodeBlock* Ctx__enter_block(Ctx* self, CodeBlockType type);
+static int Ctx__get_loop(Ctx* self, bool* has_context);
+static int Ctx__enter_block(Ctx* self, CodeBlockType type);
 static void Ctx__exit_block(Ctx* self);
 static int Ctx__emit_(Ctx* self, Opcode opcode, uint16_t arg, int line);
 static int Ctx__emit_virtual(Ctx* self, Opcode opcode, uint16_t arg, int line, bool virtual);
@@ -581,9 +581,8 @@ void CompExpr__emit_(Expr* self_, Ctx* ctx) {
     Ctx__emit_(ctx, self->op0, 0, self->line);
     vtemit_(self->iter, ctx);
     Ctx__emit_(ctx, OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
-    Ctx__enter_block(ctx, CodeBlockType_FOR_LOOP);
-    int curr_iblock = ctx->curr_iblock;
-    Ctx__emit_(ctx, OP_FOR_ITER, curr_iblock, BC_KEEPLINE);
+    int block = Ctx__enter_block(ctx, CodeBlockType_FOR_LOOP);
+    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
@@ -597,7 +596,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, curr_iblock, BC_KEEPLINE);
+    Ctx__emit_(ctx, OP_LOOP_CONTINUE, block, BC_KEEPLINE);
     Ctx__exit_block(ctx);
 }
 
@@ -1107,22 +1106,28 @@ static void Ctx__dtor(Ctx* self) {
 
 static bool is_small_int(int64_t value) { return value >= INT16_MIN && value <= INT16_MAX; }
 
-static int Ctx__get_loop(Ctx* self) {
+static bool is_context_block(CodeBlock* block) {
+    return block->type >= CodeBlockType_WITH && block->type <= CodeBlockType_FINALLY;
+}
+
+static int Ctx__get_loop(Ctx* self, bool* has_context) {
     int index = self->curr_iblock;
+    *has_context = false;
     while(index >= 0) {
         CodeBlock* block = c11__at(CodeBlock, &self->co->blocks, index);
         if(block->type == CodeBlockType_FOR_LOOP) break;
         if(block->type == CodeBlockType_WHILE_LOOP) break;
+        if(is_context_block(block)) *has_context = true;
         index = block->parent;
     }
     return index;
 }
 
-static CodeBlock* Ctx__enter_block(Ctx* self, CodeBlockType type) {
+static int Ctx__enter_block(Ctx* self, CodeBlockType type) {
     CodeBlock block = {type, self->curr_iblock, self->co->codes.length, -1, -1};
     c11_vector__push(CodeBlock, &self->co->blocks, block);
     self->curr_iblock = self->co->blocks.length - 1;
-    return c11__at(CodeBlock, &self->co->blocks, self->curr_iblock);
+    return self->curr_iblock;
 }
 
 static void Ctx__exit_block(Ctx* self) {
@@ -1336,6 +1341,18 @@ 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;
@@ -1465,7 +1482,7 @@ static Error* pop_context(Compiler* self) {
     if(co->consts.length > 65530) {
         return SyntaxError(self, "maximum number of constants exceeded");
     }
-    // pre-compute LOOP_BREAK and LOOP_CONTINUE
+    // pre-compute block.end or block.end2
     for(int i = 0; i < codes->length; i++) {
         Bytecode* bc = c11__at(Bytecode, codes, i);
         if(bc->op == OP_LOOP_CONTINUE) {
@@ -1474,6 +1491,9 @@ static Error* pop_context(Compiler* self) {
         } else if(bc->op == OP_LOOP_BREAK) {
             CodeBlock* block = c11__at(CodeBlock, &ctx()->co->blocks, bc->arg);
             Bytecode__set_signed_arg(bc, (block->end2 != -1 ? block->end2 : block->end) - i);
+        } else if(bc->op == OP_FOR_ITER || bc->op == OP_FOR_ITER_YIELD_VALUE) {
+            CodeBlock* block = c11__at(CodeBlock, &ctx()->co->blocks, bc->arg);
+            Bytecode__set_signed_arg(bc, block->end - i);
         }
     }
     // pre-compute func->is_simple
@@ -1975,18 +1995,19 @@ static Error* compile_if_stmt(Compiler* self) {
 
 static Error* compile_while_loop(Compiler* self) {
     Error* err;
-    CodeBlock* block = Ctx__enter_block(ctx(), CodeBlockType_WHILE_LOOP);
+    int block = Ctx__enter_block(ctx(), CodeBlockType_WHILE_LOOP);
     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, Ctx__get_loop(ctx()), BC_KEEPLINE, true);
+    Ctx__emit_virtual(ctx(), OP_LOOP_CONTINUE, block, BC_KEEPLINE, true);
     Ctx__patch_jump(ctx(), patch);
     Ctx__exit_block(ctx());
     // optional else clause
     if(match(TK_ELSE)) {
         check(compile_block_body(self, compile_stmt));
-        block->end2 = ctx()->co->codes.length;
+        CodeBlock* p_block = c11__at(CodeBlock, &ctx()->co->blocks, block);
+        p_block->end2 = ctx()->co->codes.length;
     }
     return NULL;
 }
@@ -1998,8 +2019,8 @@ static Error* compile_for_loop(Compiler* self) {
     check(EXPR_TUPLE(self));  // [vars, iter]
     Ctx__s_emit_top(ctx());   // [vars]
     Ctx__emit_(ctx(), OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
-    CodeBlock* block = Ctx__enter_block(ctx(), CodeBlockType_FOR_LOOP);
-    Ctx__emit_(ctx(), OP_FOR_ITER, ctx()->curr_iblock, BC_KEEPLINE);
+    int block = Ctx__enter_block(ctx(), CodeBlockType_FOR_LOOP);
+    Ctx__emit_(ctx(), OP_FOR_ITER, block, BC_KEEPLINE);
     Expr* vars = Ctx__s_popx(ctx());
     bool ok = vtemit_store(vars, ctx());
     vtdelete(vars);
@@ -2008,12 +2029,13 @@ 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, Ctx__get_loop(ctx()), BC_KEEPLINE, true);
+    Ctx__emit_virtual(ctx(), OP_LOOP_CONTINUE, block, BC_KEEPLINE, true);
     Ctx__exit_block(ctx());
     // optional else clause
     if(match(TK_ELSE)) {
         check(compile_block_body(self, compile_stmt));
-        block->end2 = ctx()->co->codes.length;
+        CodeBlock* p_block = c11__at(CodeBlock, &ctx()->co->blocks, block);
+        p_block->end2 = ctx()->co->codes.length;
     }
     return NULL;
 }
@@ -2021,12 +2043,13 @@ 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);
-    Ctx__enter_block(ctx(), CodeBlockType_FOR_LOOP);
-    Ctx__emit_(ctx(), OP_FOR_ITER_YIELD_VALUE, ctx()->curr_iblock, kw_line);
-    Ctx__emit_(ctx(), OP_LOOP_CONTINUE, Ctx__get_loop(ctx()), 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);
     Ctx__exit_block(ctx());
     // StopIteration.value will be pushed onto the stack
     return NULL;
@@ -2436,19 +2459,24 @@ static Error* compile_try_except(Compiler* self) {
     int patches[8];
     int patches_length = 0;
 
-    Ctx__enter_block(ctx(), CodeBlockType_TRY_EXCEPT);
+    Ctx__enter_block(ctx(), CodeBlockType_TRY);
     Ctx__emit_(ctx(), OP_TRY_ENTER, BC_NOARG, prev()->line);
     check(compile_block_body(self, compile_stmt));
 
-    bool finally_matched = match(TK_FINALLY);
-    if(!finally_matched) {
+    bool has_finally = curr()->type == TK_FINALLY;
+    if(!has_finally) {
         patches[patches_length++] = Ctx__emit_(ctx(), OP_JUMP_FORWARD, BC_NOARG, BC_KEEPLINE);
     }
     Ctx__exit_block(ctx());
 
-    if(finally_matched) {
+    if(has_finally) {
+        consume(TK_FINALLY);
+        Ctx__emit_(ctx(), OP_BEGIN_FINALLY, BC_NOARG, prev()->line);
         // finally only, no except block
-        compile_block_body(self, compile_stmt);
+        Ctx__enter_block(ctx(), CodeBlockType_FINALLY);
+        check(compile_block_body(self, compile_stmt));
+        Ctx__exit_block(ctx());
+        Ctx__emit_(ctx(), OP_END_FINALLY, BC_NOARG, BC_KEEPLINE);
         // re-raise if needed
         Ctx__emit_(ctx(), OP_RE_RAISE, BC_NOARG, BC_KEEPLINE);
         return NULL;
@@ -2481,7 +2509,9 @@ static Error* compile_try_except(Compiler* self) {
             Ctx__emit_(ctx(), OP_PUSH_EXCEPTION, BC_NOARG, BC_KEEPLINE);
             Ctx__emit_store_name(ctx(), name_scope(self), as_name, BC_KEEPLINE);
         }
+        Ctx__enter_block(ctx(), CodeBlockType_EXCEPT);
         check(compile_block_body(self, compile_stmt));
+        Ctx__exit_block(ctx());
         Ctx__emit_(ctx(), OP_END_EXC_HANDLING, BC_NOARG, BC_KEEPLINE);
         patches[patches_length++] = Ctx__emit_(ctx(), OP_JUMP_FORWARD, BC_NOARG, BC_KEEPLINE);
         Ctx__patch_jump(ctx(), patch);
@@ -2494,7 +2524,13 @@ static Error* compile_try_except(Compiler* self) {
     for(int i = 0; i < patches_length; i++)
         Ctx__patch_jump(ctx(), patches[i]);
 
-    if(match(TK_FINALLY)) compile_block_body(self, compile_stmt);
+    if(match(TK_FINALLY)) {
+        Ctx__emit_(ctx(), OP_BEGIN_FINALLY, BC_NOARG, prev()->line);
+        Ctx__enter_block(ctx(), CodeBlockType_FINALLY);
+        check(compile_block_body(self, compile_stmt));
+        Ctx__exit_block(ctx());
+        Ctx__emit_(ctx(), OP_END_FINALLY, BC_NOARG, BC_KEEPLINE);
+    }
     // re-raise if needed
     Ctx__emit_(ctx(), OP_RE_RAISE, BC_NOARG, BC_KEEPLINE);
     return NULL;
@@ -2508,7 +2544,8 @@ static Error* compile_stmt(Compiler* self) {
     }
     advance();
     int kw_line = prev()->line;  // backup line number
-    int curr_loop_block = Ctx__get_loop(ctx());
+    bool has_context = false;
+    int curr_loop_block = Ctx__get_loop(ctx(), &has_context);
     switch(prev()->type) {
         case TK_BREAK:
             if(curr_loop_block < 0) return SyntaxError(self, "'break' outside loop");
@@ -2522,6 +2559,7 @@ 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 {
@@ -2538,6 +2576,7 @@ 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 {
@@ -2604,7 +2643,7 @@ static Error* compile_stmt(Compiler* self) {
         case TK_WITH: {
             check(EXPR(self));  // [ <expr> ]
             Ctx__s_emit_top(ctx());
-            Ctx__enter_block(ctx(), CodeBlockType_CONTEXT_MANAGER);
+            Ctx__enter_block(ctx(), CodeBlockType_WITH);
             NameExpr* as_name = NULL;
             if(match(TK_AS)) {
                 consume(TK_ID);

src/interpreter/ceval.c

@@ -88,7 +88,6 @@ FrameResult VM__run_top_frame(VM* self) {
 
 #ifndef NDEBUG
         pk_print_stack(self, frame, byte);
-        // assert(!py_checkexc(true));
 #endif
 
         switch((Opcode)byte.op) {
@@ -708,6 +707,19 @@ FrameResult VM__run_top_frame(VM* self) {
                 }
                 return RES_YIELD;
             }
+            case OP_FOR_ITER_YIELD_VALUE: {
+                int res = py_next(TOP());
+                if(res == -1) goto __ERROR;
+                if(res) {
+                    return RES_YIELD;
+                } else {
+                    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]
+                    DISPATCH_JUMP((int16_t)byte.arg);
+                }
+            }
             /////////
             case OP_LIST_APPEND: {
                 // [list, iter, value]
@@ -771,22 +783,8 @@ FrameResult VM__run_top_frame(VM* self) {
                     DISPATCH();
                 } else {
                     assert(self->last_retval.type == tp_StopIteration);
-                    int target = Frame__prepare_loop_break(frame, &self->stack);
-                    DISPATCH_JUMP_ABSOLUTE(target);
-                }
-            }
-            case OP_FOR_ITER_YIELD_VALUE: {
-                int res = py_next(TOP());
-                if(res == -1) goto __ERROR;
-                if(res) {
-                    return RES_YIELD;
-                } else {
-                    assert(self->last_retval.type == tp_StopIteration);
-                    py_ObjectRef value = py_getslot(&self->last_retval, 0);
-                    int target = Frame__prepare_loop_break(frame, &self->stack);
-                    if(py_isnil(value)) value = py_None();
-                    PUSH(value);
-                    DISPATCH_JUMP_ABSOLUTE(target);
+                    POP();  // [iter] -> []
+                    DISPATCH_JUMP((int16_t)byte.arg);
                 }
             }
             ////////
@@ -1006,6 +1004,20 @@ FrameResult VM__run_top_frame(VM* self) {
                 py_clearexc(NULL);
                 DISPATCH();
             }
+            case OP_BEGIN_FINALLY: {
+                if(self->curr_exception.type) {
+                    // temporarily handle the exception if any
+                    self->is_curr_exc_handled = true;
+                }
+                DISPATCH();
+            }
+            case OP_END_FINALLY: {
+                if(self->curr_exception.type && self->is_curr_exc_handled) {
+                    // revert the exception handling if needed
+                    self->is_curr_exc_handled = false;
+                }
+                DISPATCH();
+            }
             //////////////////
             case OP_FORMAT_STRING: {
                 py_Ref spec = c11__at(py_TValue, &frame->co->consts, byte.arg);
@@ -1046,7 +1058,7 @@ FrameResult VM__run_top_frame(VM* self) {
     return RES_RETURN;
 }
 
-const static char* op2str(py_Name op) {
+const char* pk_op2str(py_Name op) {
     switch(op) {
         case __eq__: return "==";
         case __ne__: return "!=";
@@ -1104,7 +1116,7 @@ bool pk_stack_binaryop(VM* self, py_Name op, py_Name rop) {
         py_newbool(py_retval(), !res);
         return true;
     }
-    return TypeError("unsupported operand type(s) for '%s'", op2str(op));
+    return TypeError("unsupported operand type(s) for '%s'", pk_op2str(op));
 }
 
 bool py_binaryop(py_Ref lhs, py_Ref rhs, py_Name op, py_Name rop) {

src/interpreter/frame.c

@@ -70,7 +70,7 @@ int Frame__prepare_jump_exception_handler(Frame* self, ValueStack* _s) {
     int iblock = Frame__iblock(self);
     while(iblock >= 0) {
         CodeBlock* block = c11__at(CodeBlock, &self->co->blocks, iblock);
-        if(block->type == CodeBlockType_TRY_EXCEPT) break;
+        if(block->type == CodeBlockType_TRY) break;
         iblock = block->parent;
     }
     if(iblock < 0) return -1;
@@ -97,18 +97,11 @@ void Frame__prepare_jump_break(Frame* self, ValueStack* _s, int target) {
     }
 }
 
-int Frame__prepare_loop_break(Frame* self, ValueStack* _s) {
-    int iblock = Frame__iblock(self);
-    int target = c11__getitem(CodeBlock, &self->co->blocks, iblock).end;
-    Frame__prepare_jump_break(self, _s, target);
-    return target;
-}
-
 int Frame__exit_block(Frame* self, ValueStack* _s, int iblock) {
     CodeBlock* block = c11__at(CodeBlock, &self->co->blocks, iblock);
     if(block->type == CodeBlockType_FOR_LOOP) {
         _s->sp--;  // pop iterator
-    } else if(block->type == CodeBlockType_CONTEXT_MANAGER) {
+    } else if(block->type == CodeBlockType_WITH) {
         _s->sp--;  // pop context variable
     }
     return block->parent;

src/modules/dis.c

@@ -91,7 +91,7 @@ static void disassemble(CodeObject* co) {
                 }
                 case OP_BINARY_OP: {
                     py_Name name = byte.arg & 0xFF;
-                    pk_sprintf(&ss, " (%n)", name);
+                    pk_sprintf(&ss, " (%s)", pk_op2str(name));
                     break;
                 }
             }

src2/main.c

@@ -1,3 +1,4 @@
+#include <stdbool.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <assert.h>

tests/28_exception.py

@@ -207,10 +207,30 @@ def finally_no_match():
     finally:
         ok = True
 
+ok_2 = False
 try:
     finally_no_match()
 except KeyError:
     assert ok
-    exit(0)
+    ok_2 = True
+
+assert ok_2
+
+# finally, return (SyntaxError)
+err ='''
+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
+
 
-exit(1)

tests/73_typing.py

@@ -40,7 +40,8 @@ def converter(raw, mapper: Type[T] = None, default: T = None) -> T:
     try:
         return mapper(raw)
     except:
-        return default
+        pass
+    return default
 
 raw: str = '4'
 result: int = converter(raw, mapper=int, default=0)
@@ -53,19 +54,8 @@ def converter(raw, mapper: Callable[[Any], T] = None, default: T = None) -> T:
     try:
         return mapper(raw)
     except:
-        return default
-
-# Callable[[Any], ReturnType] means a function declare like:
-# def func(arg: Any) -> ReturnType:
-#     pass
-
-# Callable[[str, int], ReturnType] means a function declare like:
-# def func(string: str, times: int) -> ReturnType:
-#     pass
-
-# Callable[..., ReturnType] means a function declare like:
-# def func(*args, **kwargs) -> ReturnType:
-#     pass
+        pass
+    return default
 
 def is_success(value) -> bool:
     return value in (0, "OK", True, "success")