py: Allow bytecode/native to put iter_buf on stack for simple for loops.

So that the "for x in it: ..." statement can now work without using the
heap (so long as the iterator argument fits in an iter_buf structure).

1 files changed, 13 insertions, 4 deletions

diff --git a/py/emitbc.c b/py/emitbc.c
index e3a047f27..3e0c0b397 100644
--- a/py/emitbc.c
+++ b/py/emitbc.c
@@ -734,6 +734,10 @@ void mp_emit_bc_unwind_jump(emit_t *emit, mp_uint_t label, mp_uint_t except_dept
         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);
+            // also pop the iter_buf
+            for (size_t i = 0; i < sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t); ++i) {
+                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 {
@@ -773,9 +777,9 @@ void mp_emit_bc_end_finally(emit_t *emit) {
     emit_write_bytecode_byte(emit, MP_BC_END_FINALLY);
 }
 
-void mp_emit_bc_get_iter(emit_t *emit) {
-    emit_bc_pre(emit, 0);
-    emit_write_bytecode_byte(emit, MP_BC_GET_ITER);
+void mp_emit_bc_get_iter(emit_t *emit, bool use_stack) {
+    emit_bc_pre(emit, use_stack ? sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) : 0);
+    emit_write_bytecode_byte(emit, use_stack ? MP_BC_GET_ITER_STACK : MP_BC_GET_ITER);
 }
 
 void mp_emit_bc_for_iter(emit_t *emit, mp_uint_t label) {
@@ -783,8 +787,13 @@ void mp_emit_bc_for_iter(emit_t *emit, mp_uint_t label) {
     emit_write_bytecode_byte_unsigned_label(emit, MP_BC_FOR_ITER, label);
 }
 
-void mp_emit_bc_for_iter_end(emit_t *emit) {
+void mp_emit_bc_for_iter_end(emit_t *emit, bool use_stack) {
     emit_bc_pre(emit, -1);
+    if (use_stack) {
+        for (size_t i = 0; i < sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t); ++i) {
+            mp_emit_bc_pop_top(emit);
+        }
+    }
 }
 
 void mp_emit_bc_pop_block(emit_t *emit) {