py/compile: Fix bug with break/continue in else of optimised for-range.

dpgeorge · dpgeorge · commit 4c5f108321a8 · 2017-06-22T13:50:33.000+10:00
This patch fixes a bug whereby the Python stack was not correctly reset if
there was a break/continue statement in the else black of an optimised
for-range loop.

For example, in the following code the "j" variable from the inner for loop
was not being popped off the Python stack:

    for i in range(4):
        for j in range(4):
            pass
        else:
            continue

This is now fixed with this patch.
diff --git a/py/compile.c b/py/compile.c
@@ -1406,7 +1406,20 @@ STATIC void compile_for_stmt_optimised_range(compiler_t *comp, mp_parse_node_t p
     // break/continue apply to outer loop (if any) in the else block
     END_BREAK_CONTINUE_BLOCK
 
-    compile_node(comp, pn_else);
+    // Compile the else block.  We must pop the iterator variables before
+    // executing the else code because it may contain break/continue statements.
+    uint end_label = 0;
+    if (!MP_PARSE_NODE_IS_NULL(pn_else)) {
+        // discard final value of "var", and possible "end" value
+        EMIT(pop_top);
+        if (end_on_stack) {
+            EMIT(pop_top);
+        }
+        compile_node(comp, pn_else);
+        end_label = comp_next_label(comp);
+        EMIT_ARG(jump, end_label);
+        EMIT_ARG(adjust_stack_size, 1 + end_on_stack);
+    }
 
     EMIT_ARG(label_assign, break_label);
 
@@ -1417,6 +1430,10 @@ STATIC void compile_for_stmt_optimised_range(compiler_t *comp, mp_parse_node_t p
     if (end_on_stack) {
         EMIT(pop_top);
     }
+
+    if (!MP_PARSE_NODE_IS_NULL(pn_else)) {
+        EMIT_ARG(label_assign, end_label);
+    }
 }
 
 STATIC void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
@@ -1496,7 +1513,7 @@ STATIC void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     // break/continue apply to outer loop (if any) in the else block
     END_BREAK_CONTINUE_BLOCK
 
-    compile_node(comp, pns->nodes[3]); // else (not tested)
+    compile_node(comp, pns->nodes[3]); // else (may be empty)
 
     EMIT_ARG(label_assign, break_label);
 }
