diff --git a/py/objgenerator.c b/py/objgenerator.c index 543685fac7..1ee7b8b1db 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 261670d4f3..b53711bbe2 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 5365014fcb..bbfc9914eb 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);