py: Fix break from within a for loop.

Needed to pop the iterator object when breaking out of a for loop.  Need
also to be careful to unwind exception handler before popping iterator.

Addresses issue #635.

5 files changed, 20 insertions, 9 deletions

diff --git a/py/compile.c b/py/compile.c
index 31ecbf0b91..c90772a7e3 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -73,8 +73,8 @@ typedef struct _compiler_t {
 
     uint next_label;
 
-    uint break_label;
-    uint continue_label;
+    uint16_t break_label; // highest bit set indicates we are breaking out of a for loop
+    uint16_t continue_label;
     int break_continue_except_level;
     uint16_t cur_except_level; // increased for SETUP_EXCEPT, SETUP_FINALLY; decreased for POP_BLOCK, POP_EXCEPT
 
@@ -1745,6 +1745,7 @@ void compile_while_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 // And, if the loop never runs, the loop variable should never be assigned
 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) {
     START_BREAK_CONTINUE_BLOCK
+    comp->break_label |= MP_EMIT_BREAK_FROM_FOR;
 
     uint top_label = comp_next_label(comp);
     uint entry_label = comp_next_label(comp);
@@ -1843,6 +1844,7 @@ void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 #endif
 
     START_BREAK_CONTINUE_BLOCK
+    comp->break_label |= MP_EMIT_BREAK_FROM_FOR;
 
     uint pop_label = comp_next_label(comp);
     uint end_label = comp_next_label(comp);
diff --git a/py/emit.h b/py/emit.h
index 5a3b27d839..874ec8819a 100644
--- a/py/emit.h
+++ b/py/emit.h
@@ -44,6 +44,8 @@ typedef enum {
 #define MP_EMIT_STAR_FLAG_SINGLE (0x01)
 #define MP_EMIT_STAR_FLAG_DOUBLE (0x02)
 
+#define MP_EMIT_BREAK_FROM_FOR (0x8000)
+
 typedef struct _emit_t emit_t;
 
 typedef struct _emit_method_table_t {
diff --git a/py/emitbc.c b/py/emitbc.c
index 06f63b6f6c..cfaea7c88a 100644
--- a/py/emitbc.c
+++ b/py/emitbc.c
@@ -617,11 +617,15 @@ STATIC void emit_bc_jump_if_false_or_pop(emit_t *emit, uint label) {
 
 STATIC void emit_bc_unwind_jump(emit_t *emit, uint label, int except_depth) {
     if (except_depth == 0) {
-        emit_bc_jump(emit, label);
-    } else {
         emit_bc_pre(emit, 0);
-        emit_write_bytecode_byte_signed_label(emit, MP_BC_UNWIND_JUMP, label);
-        emit_write_bytecode_byte(emit, except_depth);
+        if (label & MP_EMIT_BREAK_FROM_FOR) {
+            // need to pop the iterator if we are breaking out of a for loop
+            emit_write_bytecode_byte(emit, MP_BC_POP_TOP);
+        }
+        emit_write_bytecode_byte_signed_label(emit, MP_BC_JUMP, label & ~MP_EMIT_BREAK_FROM_FOR);
+    } else {
+        emit_write_bytecode_byte_signed_label(emit, MP_BC_UNWIND_JUMP, label & ~MP_EMIT_BREAK_FROM_FOR);
+        emit_write_bytecode_byte(emit, ((label & MP_EMIT_BREAK_FROM_FOR) ? 0x80 : 0) | except_depth);
     }
 }
 
diff --git a/py/emitnative.c b/py/emitnative.c
index 057e42c756..4dac5ffb09 100644
--- a/py/emitnative.c
+++ b/py/emitnative.c
@@ -1043,7 +1043,7 @@ STATIC void emit_native_jump_if_false_or_pop(emit_t *emit, uint label) {
 }
 
 STATIC void emit_native_break_loop(emit_t *emit, uint label, int except_depth) {
-    emit_native_jump(emit, label); // TODO properly
+    emit_native_jump(emit, label & ~MP_EMIT_BREAK_FROM_FOR); // TODO properly
 }
 
 STATIC void emit_native_continue_loop(emit_t *emit, uint label, int except_depth) {
diff --git a/py/vm.c b/py/vm.c
index bd94ade541..74f821e9ec 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -618,10 +618,10 @@ dispatch_loop:
                 ENTRY(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
+                    PUSH((void*)(machine_uint_t)(*ip)); // push number of exception handlers to unwind (0x80 bit set if we also need to pop stack)
 unwind_jump:
                     unum = (machine_uint_t)POP(); // get number of exception handlers to unwind
-                    while (unum > 0) {
+                    while ((unum & 0x7f) > 0) {
                         unum -= 1;
                         assert(exc_sp >= exc_stack);
                         if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
@@ -638,6 +638,9 @@ unwind_jump:
                         exc_sp--;
                     }
                     ip = (const byte*)POP(); // pop destination ip for jump
+                    if (unum != 0) {
+                        sp--;
+                    }
                     DISPATCH();
 
                 // matched against: POP_BLOCK or POP_EXCEPT (anything else?)