py/vm: Simplify handling of MP_OBJ_STOP_ITERATION in yield-from opcode.

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

3 files changed, 10 insertions, 15 deletions

diff --git a/py/objgenerator.c b/py/objgenerator.c
index 543685fac..1ee7b8b1d 100644
--- a/py/objgenerator.c
+++ b/py/objgenerator.c
@@ -152,8 +152,9 @@ mp_vm_return_kind_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_
     mp_check_self(mp_obj_is_type(self_in, &mp_type_gen_instance));
     mp_obj_gen_instance_t *self = MP_OBJ_TO_PTR(self_in);
     if (self->code_state.ip == 0) {
-        // Trying to resume already stopped generator
-        *ret_val = MP_OBJ_STOP_ITERATION;
+        // Trying to resume an already stopped generator.
+        // This is an optimised "raise StopIteration(None)".
+        *ret_val = mp_const_none;
         return MP_VM_RETURN_NORMAL;
     }
 
@@ -212,6 +213,7 @@ mp_vm_return_kind_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_
             // subsequent next() may re-execute statements after last yield
             // again and again, leading to side effects.
             self->code_state.ip = 0;
+            // This is an optimised "raise StopIteration(*ret_val)".
             *ret_val = *self->code_state.sp;
             break;
 
diff --git a/py/runtime.c b/py/runtime.c
index 261670d4f..b53711bbe 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -1272,7 +1272,6 @@ mp_obj_t mp_iternext(mp_obj_t o_in) {
     }
 }
 
-// TODO: Unclear what to do with StopIterarion exception here.
 mp_vm_return_kind_t mp_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value, mp_obj_t *ret_val) {
     assert((send_value != MP_OBJ_NULL) ^ (throw_value != MP_OBJ_NULL));
     const mp_obj_type_t *type = mp_obj_get_type(self_in);
@@ -1287,8 +1286,9 @@ mp_vm_return_kind_t mp_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t th
         if (ret != MP_OBJ_STOP_ITERATION) {
             return MP_VM_RETURN_YIELD;
         } else {
-            // Emulate raise StopIteration()
-            // Special case, handled in vm.c
+            // The generator is finished.
+            // This is an optimised "raise StopIteration(None)".
+            *ret_val = mp_const_none;
             return MP_VM_RETURN_NORMAL;
         }
     }
diff --git a/py/vm.c b/py/vm.c
index 5365014fc..bbfc9914e 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -1257,16 +1257,9 @@ yield:
                         PUSH(ret_value);
                         goto yield;
                     } else if (ret_kind == MP_VM_RETURN_NORMAL) {
-                        // Pop exhausted gen
-                        sp--;
-                        if (ret_value == MP_OBJ_STOP_ITERATION) {
-                            // Optimize StopIteration
-                            // TODO: get StopIteration's value
-                            PUSH(mp_const_none);
-                        } else {
-                            PUSH(ret_value);
-                        }
-
+                        // The generator has finished, and returned a value via StopIteration
+                        // Replace exhausted generator with the returned value
+                        SET_TOP(ret_value);
                         // If we injected GeneratorExit downstream, then even
                         // if it was swallowed, we re-raise GeneratorExit
                         GENERATOR_EXIT_IF_NEEDED(t_exc);