py: Support single argument to optimised MP_OBJ_STOP_ITERATION.

The MP_OBJ_STOP_ITERATION optimisation is a shortcut for creating a
StopIteration() exception object, and means that heap memory does not need
to be allocated for the exception (in cases where it can be used).  This
commit allows this optimised object to take an optional argument (before,
it could only have no argument).

The commit also adds some new tests to cover corner cases with
StopIteration and generators that previously did not work.

Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 17 insertions, 3 deletions

diff --git a/py/runtime.c b/py/runtime.c
index b53711bbe..2c849fe95 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -1217,6 +1217,7 @@ mp_obj_t mp_getiter(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) {
 mp_obj_t mp_iternext_allow_raise(mp_obj_t o_in) {
     const mp_obj_type_t *type = mp_obj_get_type(o_in);
     if (type->iternext != NULL) {
+        MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL;
         return type->iternext(o_in);
     } else {
         // check for __next__ method
@@ -1242,6 +1243,7 @@ mp_obj_t mp_iternext(mp_obj_t o_in) {
     MP_STACK_CHECK(); // enumerate, filter, map and zip can recursively call mp_iternext
     const mp_obj_type_t *type = mp_obj_get_type(o_in);
     if (type->iternext != NULL) {
+        MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL;
         return type->iternext(o_in);
     } else {
         // check for __next__ method
@@ -1256,7 +1258,7 @@ mp_obj_t mp_iternext(mp_obj_t o_in) {
                 return ret;
             } else {
                 if (mp_obj_is_subclass_fast(MP_OBJ_FROM_PTR(((mp_obj_base_t *)nlr.ret_val)->type), MP_OBJ_FROM_PTR(&mp_type_StopIteration))) {
-                    return MP_OBJ_STOP_ITERATION;
+                    return mp_make_stop_iteration(mp_obj_exception_get_value(MP_OBJ_FROM_PTR(nlr.ret_val)));
                 } else {
                     nlr_jump(nlr.ret_val);
                 }
@@ -1281,14 +1283,18 @@ mp_vm_return_kind_t mp_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t th
     }
 
     if (type->iternext != NULL && send_value == mp_const_none) {
+        MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL;
         mp_obj_t ret = type->iternext(self_in);
         *ret_val = ret;
         if (ret != MP_OBJ_STOP_ITERATION) {
             return MP_VM_RETURN_YIELD;
         } else {
             // The generator is finished.
-            // This is an optimised "raise StopIteration(None)".
-            *ret_val = mp_const_none;
+            // This is an optimised "raise StopIteration(*ret_val)".
+            *ret_val = MP_STATE_THREAD(stop_iteration_arg);
+            if (*ret_val == MP_OBJ_NULL) {
+                *ret_val = mp_const_none;
+            }
             return MP_VM_RETURN_NORMAL;
         }
     }
@@ -1559,6 +1565,14 @@ NORETURN void mp_raise_NotImplementedError(mp_rom_error_text_t msg) {
 
 #endif
 
+NORETURN void mp_raise_StopIteration(mp_obj_t arg) {
+    if (arg == MP_OBJ_NULL) {
+        mp_raise_type(&mp_type_StopIteration);
+    } else {
+        nlr_raise(mp_obj_new_exception_arg1(&mp_type_StopIteration, arg));
+    }
+}
+
 NORETURN void mp_raise_OSError(int errno_) {
     nlr_raise(mp_obj_new_exception_arg1(&mp_type_OSError, MP_OBJ_NEW_SMALL_INT(errno_)));
 }