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, 11 insertions, 0 deletions

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