py: Implement break/continue from an exception with finally.

Still todo: break/continue from within the finally block itself.

py/bc0.h

@@ -52,8 +52,6 @@
 #define MP_BC_JUMP_IF_TRUE_OR_POP    (0x48) // rel byte code offset, 16-bit signed, in excess
 #define MP_BC_JUMP_IF_FALSE_OR_POP   (0x49) // rel byte code offset, 16-bit signed, in excess
 #define MP_BC_SETUP_LOOP         (0x4a) // rel byte code offset, 16-bit unsigned
-#define MP_BC_BREAK_LOOP         (0x4b) // rel byte code offset, 16-bit unsigned
-#define MP_BC_CONTINUE_LOOP      (0x4c) // rel byte code offset, 16-bit unsigned
 #define MP_BC_SETUP_WITH         (0x4d) // rel byte code offset, 16-bit unsigned
 #define MP_BC_WITH_CLEANUP       (0x4e)
 #define MP_BC_SETUP_EXCEPT       (0x4f) // rel byte code offset, 16-bit unsigned
@@ -63,6 +61,7 @@
 #define MP_BC_FOR_ITER           (0x53) // rel byte code offset, 16-bit unsigned
 #define MP_BC_POP_BLOCK          (0x54)
 #define MP_BC_POP_EXCEPT         (0x55)
+#define MP_BC_UNWIND_JUMP        (0x56) // rel byte code offset, 16-bit signed, in excess; then a byte
 
 #define MP_BC_UNARY_OP           (0x60) // byte
 #define MP_BC_BINARY_OP          (0x61) // byte

py/compile.c

@@ -50,7 +50,8 @@ typedef struct _compiler_t {
 
     int break_label;
     int continue_label;
-    int except_nest_level;
+    int break_continue_except_level;
+    int cur_except_level; // increased for SETUP_EXCEPT, SETUP_FINALLY; decreased for POP_BLOCK, POP_EXCEPT
 
     int n_arg_keyword;
     bool have_star_arg;
@@ -1080,18 +1081,14 @@ void compile_break_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     if (comp->break_label == 0) {
         printf("ERROR: cannot break from here\n");
     }
-    EMIT_ARG(break_loop, comp->break_label);
+    EMIT_ARG(break_loop, comp->break_label, comp->cur_except_level - comp->break_continue_except_level);
 }
 
 void compile_continue_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     if (comp->continue_label == 0) {
         printf("ERROR: cannot continue from here\n");
     }
-    if (comp->except_nest_level > 0) {
-        EMIT_ARG(continue_loop, comp->continue_label);
-    } else {
-        EMIT_ARG(jump, comp->continue_label);
-    }
+    EMIT_ARG(continue_loop, comp->continue_label, comp->cur_except_level - comp->break_continue_except_level);
 }
 
 void compile_return_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
@@ -1387,15 +1384,22 @@ void compile_if_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     EMIT_ARG(label_assign, l_end);
 }
 
-void compile_while_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
-    int old_break_label = comp->break_label;
-    int old_continue_label = comp->continue_label;
+#define START_BREAK_CONTINUE_BLOCK \
+    int old_break_label = comp->break_label; \
+    int old_continue_label = comp->continue_label; \
+    int break_label = comp_next_label(comp); \
+    int continue_label = comp_next_label(comp); \
+    comp->break_label = break_label; \
+    comp->continue_label = continue_label; \
+    comp->break_continue_except_level = comp->cur_except_level;
 
-    int break_label = comp_next_label(comp);
-    int continue_label = comp_next_label(comp);
+#define END_BREAK_CONTINUE_BLOCK \
+    comp->break_label = old_break_label; \
+    comp->continue_label = old_continue_label; \
+    comp->break_continue_except_level = comp->cur_except_level;
 
-    comp->break_label = break_label;
-    comp->continue_label = continue_label;
+void compile_while_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
+    START_BREAK_CONTINUE_BLOCK
 
     // compared to CPython, we have an optimised version of while loops
 #if MICROPY_EMIT_CPYTHON
@@ -1423,8 +1427,7 @@ void compile_while_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 #endif
 
     // break/continue apply to outer loop (if any) in the else block
-    comp->break_label = old_break_label;
-    comp->continue_label = old_continue_label;
+    END_BREAK_CONTINUE_BLOCK
 
     compile_node(comp, pns->nodes[2]); // else
 
@@ -1434,14 +1437,7 @@ void compile_while_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 // TODO preload end and step onto stack if they are not constants
 // TODO check if step is negative and do opposite test
 void compile_for_stmt_optimised_range(compiler_t *comp, mp_parse_node_t pn_var, mp_parse_node_t pn_start, mp_parse_node_t pn_end, mp_parse_node_t pn_step, mp_parse_node_t pn_body, mp_parse_node_t pn_else) {
-    int old_break_label = comp->break_label;
-    int old_continue_label = comp->continue_label;
-
-    int break_label = comp_next_label(comp);
-    int continue_label = comp_next_label(comp);
-
-    comp->break_label = break_label;
-    comp->continue_label = continue_label;
+    START_BREAK_CONTINUE_BLOCK
 
     int top_label = comp_next_label(comp);
     int entry_label = comp_next_label(comp);
@@ -1477,8 +1473,7 @@ void compile_for_stmt_optimised_range(compiler_t *comp, mp_parse_node_t pn_var,
     EMIT_ARG(pop_jump_if_true, top_label);
 
     // break/continue apply to outer loop (if any) in the else block
-    comp->break_label = old_break_label;
-    comp->continue_label = old_continue_label;
+    END_BREAK_CONTINUE_BLOCK
 
     compile_node(comp, pn_else);
 
@@ -1531,18 +1526,11 @@ void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     }
 #endif
 
-    int old_break_label = comp->break_label;
-    int old_continue_label = comp->continue_label;
+    START_BREAK_CONTINUE_BLOCK
 
-    int for_label = comp_next_label(comp);
     int pop_label = comp_next_label(comp);
     int end_label = comp_next_label(comp);
 
-    int break_label = comp_next_label(comp);
-
-    comp->continue_label = for_label;
-    comp->break_label = break_label;
-
     // I don't think our implementation needs SETUP_LOOP/POP_BLOCK for for-statements
 #if MICROPY_EMIT_CPYTHON
     EMIT_ARG(setup_loop, end_label);
@@ -1550,19 +1538,18 @@ void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 
     compile_node(comp, pns->nodes[1]); // iterator
     EMIT(get_iter);
-    EMIT_ARG(label_assign, for_label);
+    EMIT_ARG(label_assign, continue_label);
     EMIT_ARG(for_iter, pop_label);
     c_assign(comp, pns->nodes[0], ASSIGN_STORE); // variable
     compile_node(comp, pns->nodes[2]); // body
     if (!EMIT(last_emit_was_return_value)) {
-        EMIT_ARG(jump, for_label);
+        EMIT_ARG(jump, continue_label);
     }
     EMIT_ARG(label_assign, pop_label);
     EMIT(for_iter_end);
 
     // break/continue apply to outer loop (if any) in the else block
-    comp->break_label = old_break_label;
-    comp->continue_label = old_continue_label;
+    END_BREAK_CONTINUE_BLOCK
 
 #if MICROPY_EMIT_CPYTHON
     EMIT(pop_block);
@@ -1582,8 +1569,10 @@ void compile_try_except(compiler_t *comp, mp_parse_node_t pn_body, int n_except,
     int stack_size = EMIT(get_stack_size);
     int l1 = comp_next_label(comp);
     int success_label = comp_next_label(comp);
-    comp->except_nest_level += 1; // for correct handling of continue
+
     EMIT_ARG(setup_except, l1);
+    comp->cur_except_level += 1;
+
     compile_node(comp, pn_body); // body
     EMIT(pop_block);
     EMIT_ARG(jump, success_label);
@@ -1634,6 +1623,7 @@ void compile_try_except(compiler_t *comp, mp_parse_node_t pn_body, int n_except,
         if (qstr_exception_local != 0) {
             l3 = comp_next_label(comp);
             EMIT_ARG(setup_finally, l3);
+            comp->cur_except_level += 1;
         }
         compile_node(comp, pns_except->nodes[1]);
         if (qstr_exception_local != 0) {
@@ -1646,15 +1636,18 @@ void compile_try_except(compiler_t *comp, mp_parse_node_t pn_body, int n_except,
             EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
             EMIT_ARG(store_id, qstr_exception_local);
             EMIT_ARG(delete_id, qstr_exception_local);
+
+            comp->cur_except_level -= 1;
             EMIT(end_finally);
         }
         EMIT_ARG(jump, l2);
         EMIT_ARG(label_assign, end_finally_label);
     }
 
+    comp->cur_except_level -= 1;
     EMIT(end_finally);
+
     EMIT_ARG(label_assign, success_label);
-    comp->except_nest_level -= 1;
     compile_node(comp, pn_else); // else block, can be null
     EMIT_ARG(label_assign, l2);
     EMIT_ARG(set_stack_size, stack_size);
@@ -1664,7 +1657,10 @@ void compile_try_finally(compiler_t *comp, mp_parse_node_t pn_body, int n_except
     // don't understand how the stack works with exceptions, so we force it to return to the correct value
     int stack_size = EMIT(get_stack_size);
     int l_finally_block = comp_next_label(comp);
+
     EMIT_ARG(setup_finally, l_finally_block);
+    comp->cur_except_level += 1;
+
     if (n_except == 0) {
         assert(MP_PARSE_NODE_IS_NULL(pn_else));
         compile_node(comp, pn_body);
@@ -1675,7 +1671,10 @@ void compile_try_finally(compiler_t *comp, mp_parse_node_t pn_body, int n_except
     EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
     EMIT_ARG(label_assign, l_finally_block);
     compile_node(comp, pn_finally);
+
+    comp->cur_except_level -= 1;
     EMIT(end_finally);
+
     EMIT_ARG(set_stack_size, stack_size);
 }
 
@@ -3056,7 +3055,9 @@ mp_obj_t mp_compile(mp_parse_node_t pn, qstr source_file, bool is_repl) {
 
     comp->break_label = 0;
     comp->continue_label = 0;
-    comp->except_nest_level = 0;
+    comp->break_continue_except_level = 0;
+    comp->cur_except_level = 0;
+
     comp->scope_head = NULL;
     comp->scope_cur = NULL;
 

py/emit.h

@@ -72,8 +72,8 @@ typedef struct _emit_method_table_t {
     void (*jump_if_true_or_pop)(emit_t *emit, int label);
     void (*jump_if_false_or_pop)(emit_t *emit, int label);
     void (*setup_loop)(emit_t *emit, int label);
-    void (*break_loop)(emit_t *emit, int label);
-    void (*continue_loop)(emit_t *emit, int label);
+    void (*break_loop)(emit_t *emit, int label, int except_depth);
+    void (*continue_loop)(emit_t *emit, int label, int except_depth);
     void (*setup_with)(emit_t *emit, int label);
     void (*with_cleanup)(emit_t *emit);
     void (*setup_except)(emit_t *emit, int label);

py/emitbc.c

@@ -540,14 +540,14 @@ static void emit_bc_setup_loop(emit_t *emit, int label) {
     emit_write_byte_code_byte_unsigned_label(emit, MP_BC_SETUP_LOOP, label);
 }
 
-static void emit_bc_break_loop(emit_t *emit, int label) {
-    emit_pre(emit, 0);
-    emit_write_byte_code_byte_unsigned_label(emit, MP_BC_BREAK_LOOP, label);
-}
-
-static void emit_bc_continue_loop(emit_t *emit, int label) {
-    emit_pre(emit, 0);
-    emit_write_byte_code_byte_unsigned_label(emit, MP_BC_CONTINUE_LOOP, label);
+static void emit_bc_unwind_jump(emit_t *emit, int label, int except_depth) {
+    if (except_depth == 0) {
+        emit_bc_jump(emit, label);
+    } else {
+        emit_pre(emit, 0);
+        emit_write_byte_code_byte_signed_label(emit, MP_BC_UNWIND_JUMP, label);
+        emit_write_byte_code_byte(emit, except_depth);
+    }
 }
 
 static void emit_bc_setup_with(emit_t *emit, int label) {
@@ -828,8 +828,8 @@ const emit_method_table_t emit_bc_method_table = {
     emit_bc_jump_if_true_or_pop,
     emit_bc_jump_if_false_or_pop,
     emit_bc_setup_loop,
-    emit_bc_break_loop,
-    emit_bc_continue_loop,
+    emit_bc_unwind_jump,
+    emit_bc_unwind_jump,
     emit_bc_setup_with,
     emit_bc_with_cleanup,
     emit_bc_setup_except,

py/emitnative.c

@@ -931,10 +931,10 @@ static void emit_native_setup_loop(emit_t *emit, int label) {
     emit_post(emit);
 }
 
-static void emit_native_break_loop(emit_t *emit, int label) {
+static void emit_native_break_loop(emit_t *emit, int label, int except_depth) {
     emit_native_jump(emit, label); // TODO properly
 }
-static void emit_native_continue_loop(emit_t *emit, int label) {
+static void emit_native_continue_loop(emit_t *emit, int label, int except_depth) {
     assert(0);
 }
 static void emit_native_setup_with(emit_t *emit, int label) {

py/showbc.c

@@ -220,14 +220,10 @@ void mp_byte_code_print(const byte *ip, int len) {
                 printf("SETUP_LOOP " UINT_FMT, ip + unum - ip_start);
                 break;
 
-            case MP_BC_BREAK_LOOP:
-                DECODE_ULABEL; // loop labels are always forward
-                printf("BREAK_LOOP " UINT_FMT, ip + unum - ip_start);
-                break;
-
-            case MP_BC_CONTINUE_LOOP:
-                DECODE_ULABEL; // loop labels are always forward
-                printf("CONTINUE_LOOP " UINT_FMT, ip + unum - ip_start);
+            case MP_BC_UNWIND_JUMP:
+                DECODE_SLABEL;
+                printf("UNWIND_JUMP " UINT_FMT " %d", ip + unum - ip_start, *ip);
+                ip += 1;
                 break;
 
             case MP_BC_SETUP_EXCEPT:

py/vm.c

@@ -31,10 +31,11 @@ typedef struct _mp_exc_stack {
 } mp_exc_stack;
 
 // Exception stack unwind reasons (WHY_* in CPython-speak)
+// TODO perhaps compress this to RETURN=0, JUMP>0, with number of unwinds
+// left to do encoded in the JUMP number
 typedef enum {
     UNWIND_RETURN = 1,
-    UNWIND_BREAK,
-    UNWIND_CONTINUE,
+    UNWIND_JUMP,
 } mp_unwind_reason_t;
 
 #define DECODE_UINT do { unum = *ip++; if (unum > 127) { unum = ((unum & 0x3f) << 8) | (*ip++); } } while (0)
@@ -328,16 +329,29 @@ dispatch_loop:
                         break;
                         */
 
-                    // TODO this might need more sophisticated handling when breaking from within an except
-                    case MP_BC_BREAK_LOOP:
-                        DECODE_ULABEL;
-                        ip += unum;
-                        break;
-
-                    // TODO this might need more sophisticated handling when breaking from within an except
-                    case MP_BC_CONTINUE_LOOP:
-                        DECODE_ULABEL;
-                        ip += unum;
+                    case MP_BC_UNWIND_JUMP:
+                        DECODE_SLABEL;
+                        PUSH((void*)(ip + unum)); // push destination ip for jump
+                        PUSH((void*)(machine_uint_t)(*ip)); // push number of exception handlers to unwind
+unwind_jump:
+                        unum = (machine_uint_t)POP(); // get number of exception handlers to unwind
+                        while (unum > 0) {
+                            unum -= 1;
+                            assert(exc_sp >= exc_stack);
+                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY) {
+                                // We're going to run "finally" code as a coroutine
+                                // (not calling it recursively). Set up a sentinel
+                                // on a stack so it can return back to us when it is
+                                // done (when END_FINALLY reached).
+                                PUSH((void*)unum); // push number of exception handlers left to unwind
+                                PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_JUMP)); // push sentinel
+                                ip = exc_sp->handler; // get exception handler byte code address
+                                exc_sp--; // pop exception handler
+                                goto dispatch_loop; // run the exception handler
+                            }
+                            exc_sp--;
+                        }
+                        ip = (const byte*)POP(); // pop destination ip for jump
                         break;
 
                     // matched against: POP_BLOCK or POP_EXCEPT (anything else?)
@@ -369,10 +383,8 @@ dispatch_loop:
                             switch (reason) {
                                 case UNWIND_RETURN:
                                     goto unwind_return;
-                                // TODO
-                                case UNWIND_BREAK:
-                                case UNWIND_CONTINUE:
-                                ;
+                                case UNWIND_JUMP:
+                                    goto unwind_jump;
                             }
                             assert(0);
                         } else {

tests/basics/try-finally-break.py

+for i in range(4):
+    print(i)
+    try:
+        while True:
+            try:
+                try:
+                    break
+                finally:
+                    print('finally 1')
+            finally:
+                print('finally 2')
+        print('here')
+    finally:
+        print('finnaly 3')