py: Fix stack underflow with optimised for loop.

dpgeorge · dpgeorge · commit 5b5562c1d13f · 2014-05-31T17:59:11.000+01:00
diff --git a/py/compile.c b/py/compile.c
@@ -1745,7 +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;
+    // note that we don't need to pop anything when breaking from an optimise for loop
 
     uint top_label = comp_next_label(comp);
     uint entry_label = comp_next_label(comp);
diff --git a/py/vm.c b/py/vm.c
@@ -159,8 +159,8 @@ mp_vm_return_kind_t mp_execute_bytecode(const byte *code, const mp_obj_t *args,
 
 #if DETECT_VM_STACK_OVERFLOW
     if (vm_return_kind == MP_VM_RETURN_NORMAL) {
-        if (sp != state) {
-            printf("Stack misalign: %d\n", sp - state);
+        if (sp < state) {
+            printf("VM stack underflow: " INT_FMT "\n", sp - state);
             assert(0);
         }
     }
@@ -178,7 +178,7 @@ mp_vm_return_kind_t mp_execute_bytecode(const byte *code, const mp_obj_t *args,
             }
         }
         if (overflow) {
-            printf("VM stack overflow state=%p n_state+1=%u\n", state, n_state);
+            printf("VM stack overflow state=%p n_state+1=" UINT_FMT "\n", state, n_state);
             assert(0);
         }
     }
diff --git a/tests/basics/for_return.py b/tests/basics/for_return.py
@@ -0,0 +1,7 @@
+# test returning from within a for loop
+
+def f():
+    for i in [1, 2, 3]:
+        return i
+
+print(f())

Original file line number	Diff line number	Diff line change
`@@ -159,8 +159,8 @@ mp_vm_return_kind_t mp_execute_bytecode(const byte code, const mp_obj_t args,`
`159`	`159`
`160`	`160`	`#if DETECT_VM_STACK_OVERFLOW`
`161`	`161`	`if (vm_return_kind == MP_VM_RETURN_NORMAL) {`
`162`		`- if (sp != state) {`
`163`		`- printf("Stack misalign: %d\n", sp - state);`
	`162`	`+ if (sp < state) {`
	`163`	`+ printf("VM stack underflow: " INT_FMT "\n", sp - state);`
`164`	`164`	`assert(0);`
`165`	`165`	`}`
`166`	`166`	`}`
`@@ -178,7 +178,7 @@ mp_vm_return_kind_t mp_execute_bytecode(const byte code, const mp_obj_t args,`
`178`	`178`	`}`
`179`	`179`	`}`
`180`	`180`	`if (overflow) {`
`181`		`- printf("VM stack overflow state=%p n_state+1=%u\n", state, n_state);`
	`181`	`+ printf("VM stack overflow state=%p n_state+1=" UINT_FMT "\n", state, n_state);`
`182`	`182`	`assert(0);`
`183`	`183`	`}`
`184`	`184`	`}`