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).

8 files changed, 51 insertions, 17 deletions

diff --git a/py/bc0.h b/py/bc0.h
index 5ff9e50a89..c2b019f1a9 100644
--- a/py/bc0.h
+++ b/py/bc0.h
@@ -79,6 +79,7 @@
 #define MP_BC_POP_BLOCK          (0x44)
 #define MP_BC_POP_EXCEPT         (0x45)
 #define MP_BC_UNWIND_JUMP        (0x46) // rel byte code offset, 16-bit signed, in excess; then a byte
+#define MP_BC_GET_ITER_STACK     (0x47)
 
 #define MP_BC_BUILD_TUPLE        (0x50) // uint
 #define MP_BC_BUILD_LIST         (0x51) // uint
diff --git a/py/compile.c b/py/compile.c
index b84793d10a..d2fa03f01e 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1475,7 +1475,7 @@ STATIC void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     uint pop_label = comp_next_label(comp);
 
     compile_node(comp, pns->nodes[1]); // iterator
-    EMIT(get_iter);
+    EMIT_ARG(get_iter, true);
     EMIT_ARG(label_assign, continue_label);
     EMIT_ARG(for_iter, pop_label);
     c_assign(comp, pns->nodes[0], ASSIGN_STORE); // variable
@@ -1484,7 +1484,7 @@ STATIC void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
         EMIT_ARG(jump, continue_label);
     }
     EMIT_ARG(label_assign, pop_label);
-    EMIT(for_iter_end);
+    EMIT_ARG(for_iter_end, true);
 
     // break/continue apply to outer loop (if any) in the else block
     END_BREAK_CONTINUE_BLOCK
@@ -1680,7 +1680,7 @@ STATIC void compile_with_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 }
 
 STATIC void compile_yield_from(compiler_t *comp) {
-    EMIT(get_iter);
+    EMIT_ARG(get_iter, false);
     EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
     EMIT(yield_from);
 }
@@ -2372,7 +2372,7 @@ STATIC void compile_comprehension(compiler_t *comp, mp_parse_node_struct_t *pns,
     close_over_variables_etc(comp, this_scope, 0, 0);
 
     compile_node(comp, pns_comp_for->nodes[1]); // source of the iterator
-    EMIT(get_iter);
+    EMIT_ARG(get_iter, false);
     EMIT_ARG(call_function, 1, 0, 0);
 }
 
@@ -2900,13 +2900,13 @@ STATIC void compile_scope_comp_iter(compiler_t *comp, mp_parse_node_struct_t *pn
         // for loop
         mp_parse_node_struct_t *pns_comp_for2 = (mp_parse_node_struct_t*)pn_iter;
         compile_node(comp, pns_comp_for2->nodes[1]);
-        EMIT(get_iter);
+        EMIT_ARG(get_iter, false);
         compile_scope_comp_iter(comp, pns_comp_for2, pn_inner_expr, for_depth + 1);
     }
 
     EMIT_ARG(jump, l_top);
     EMIT_ARG(label_assign, l_end);
-    EMIT(for_iter_end);
+    EMIT_ARG(for_iter_end, false);
 }
 
 STATIC void check_for_doc_string(compiler_t *comp, mp_parse_node_t pn) {
diff --git a/py/emit.h b/py/emit.h
index 7d00f11f95..6acae9e80f 100644
--- a/py/emit.h
+++ b/py/emit.h
@@ -110,9 +110,9 @@ typedef struct _emit_method_table_t {
     void (*setup_except)(emit_t *emit, mp_uint_t label);
     void (*setup_finally)(emit_t *emit, mp_uint_t label);
     void (*end_finally)(emit_t *emit);
-    void (*get_iter)(emit_t *emit);
+    void (*get_iter)(emit_t *emit, bool use_stack);
     void (*for_iter)(emit_t *emit, mp_uint_t label);
-    void (*for_iter_end)(emit_t *emit);
+    void (*for_iter_end)(emit_t *emit, bool use_stack);
     void (*pop_block)(emit_t *emit);
     void (*pop_except)(emit_t *emit);
     void (*unary_op)(emit_t *emit, mp_unary_op_t op);
@@ -228,9 +228,9 @@ void mp_emit_bc_with_cleanup(emit_t *emit, mp_uint_t label);
 void mp_emit_bc_setup_except(emit_t *emit, mp_uint_t label);
 void mp_emit_bc_setup_finally(emit_t *emit, mp_uint_t label);
 void mp_emit_bc_end_finally(emit_t *emit);
-void mp_emit_bc_get_iter(emit_t *emit);
+void mp_emit_bc_get_iter(emit_t *emit, bool use_stack);
 void mp_emit_bc_for_iter(emit_t *emit, mp_uint_t label);
-void mp_emit_bc_for_iter_end(emit_t *emit);
+void mp_emit_bc_for_iter_end(emit_t *emit, bool use_stack);
 void mp_emit_bc_pop_block(emit_t *emit);
 void mp_emit_bc_pop_except(emit_t *emit);
 void mp_emit_bc_unary_op(emit_t *emit, mp_unary_op_t op);
diff --git a/py/emitbc.c b/py/emitbc.c
index e3a047f272..3e0c0b3972 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) {
diff --git a/py/emitnative.c b/py/emitnative.c
index e8e3754e17..6e8a37bdea 100644
--- a/py/emitnative.c
+++ b/py/emitnative.c
@@ -1799,14 +1799,19 @@ STATIC void emit_native_end_finally(emit_t *emit) {
     emit_post(emit);
 }
 
-STATIC void emit_native_get_iter(emit_t *emit) {
+STATIC void emit_native_get_iter(emit_t *emit, bool use_stack) {
     // perhaps the difficult one, as we want to rewrite for loops using native code
     // in cases where we iterate over a Python object, can we use normal runtime calls?
 
     vtype_kind_t vtype;
     emit_pre_pop_reg(emit, &vtype, REG_ARG_1);
     assert(vtype == VTYPE_PYOBJ);
-    assert(0); // TODO allocate memory for iter_buf
+    if (use_stack) {
+        emit_get_stack_pointer_to_reg_for_push(emit, REG_ARG_2, sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t));
+    } else {
+        // mp_getiter will allocate the iter_buf on the heap
+        ASM_MOV_IMM_TO_REG(emit->as, 0, REG_ARG_2);
+    }
     emit_call(emit, MP_F_GETITER);
     emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
 }
@@ -1822,10 +1827,13 @@ STATIC void emit_native_for_iter(emit_t *emit, mp_uint_t label) {
     emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
 }
 
-STATIC void emit_native_for_iter_end(emit_t *emit) {
+STATIC void emit_native_for_iter_end(emit_t *emit, bool use_stack) {
     // adjust stack counter (we get here from for_iter ending, which popped the value for us)
     emit_native_pre(emit);
     adjust_stack(emit, -1);
+    if (use_stack) {
+        adjust_stack(emit, -(sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t)));
+    }
     emit_post(emit);
 }
 
diff --git a/py/showbc.c b/py/showbc.c
index 9d20decdc8..b52905f677 100644
--- a/py/showbc.c
+++ b/py/showbc.c
@@ -387,6 +387,10 @@ const byte *mp_bytecode_print_str(const byte *ip) {
             printf("GET_ITER");
             break;
 
+        case MP_BC_GET_ITER_STACK:
+            printf("GET_ITER_STACK");
+            break;
+
         case MP_BC_FOR_ITER:
             DECODE_ULABEL; // the jump offset if iteration finishes; for labels are always forward
             printf("FOR_ITER " UINT_FMT, (mp_uint_t)(ip + unum - mp_showbc_code_start));
diff --git a/py/vm.c b/py/vm.c
index 917d41a11b..848a77a453 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -681,7 +681,9 @@ unwind_jump:;
                     }
                     ip = (const byte*)MP_OBJ_TO_PTR(POP()); // pop destination ip for jump
                     if (unum != 0) {
+                        // pop iter and iter_buf
                         sp--;
+                        sp -= sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t);
                     }
                     DISPATCH_WITH_PEND_EXC_CHECK();
                 }
@@ -726,6 +728,15 @@ unwind_jump:;
                     SET_TOP(mp_getiter(TOP(), NULL));
                     DISPATCH();
 
+                ENTRY(MP_BC_GET_ITER_STACK): {
+                    MARK_EXC_IP_SELECTIVE();
+                    mp_obj_t obj = TOP();
+                    mp_obj_iter_buf_t *iter_buf = (mp_obj_iter_buf_t*)sp;
+                    sp += sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t);
+                    SET_TOP(mp_getiter(obj, iter_buf));
+                    DISPATCH();
+                }
+
                 ENTRY(MP_BC_FOR_ITER): {
                     MARK_EXC_IP_SELECTIVE();
                     DECODE_ULABEL; // the jump offset if iteration finishes; for labels are always forward
diff --git a/py/vmentrytable.h b/py/vmentrytable.h
index dd30dd7a54..8731c3d4c4 100644
--- a/py/vmentrytable.h
+++ b/py/vmentrytable.h
@@ -73,6 +73,7 @@ static const void *const entry_table[256] = {
     [MP_BC_SETUP_FINALLY] = &&entry_MP_BC_SETUP_FINALLY,
     [MP_BC_END_FINALLY] = &&entry_MP_BC_END_FINALLY,
     [MP_BC_GET_ITER] = &&entry_MP_BC_GET_ITER,
+    [MP_BC_GET_ITER_STACK] = &&entry_MP_BC_GET_ITER_STACK,
     [MP_BC_FOR_ITER] = &&entry_MP_BC_FOR_ITER,
     [MP_BC_POP_BLOCK] = &&entry_MP_BC_POP_BLOCK,
     [MP_BC_POP_EXCEPT] = &&entry_MP_BC_POP_EXCEPT,