[mypyc] Add am_send slot and use PyIter_Send for await/yield from

mypyc/codegen/emitclass.py

@@ -14,6 +14,7 @@
 )
 from mypyc.codegen.emitfunc import native_function_header
 from mypyc.codegen.emitwrapper import (
+    generate_am_send_wrapper,
     generate_bin_op_wrapper,
     generate_bool_wrapper,
     generate_contains_wrapper,
@@ -29,6 +30,7 @@
     BITMAP_BITS,
     BITMAP_TYPE,
     CPYFUNCTION_NAME,
+    GENERATOR_HELPER_NAME,
     IS_FREE_THREADED,
     MYPYC_DEFAULTS_SETUP,
     NATIVE_PREFIX,
@@ -149,6 +151,9 @@ def dunder_attr_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     "__await__": ("am_await", native_slot),
     "__aiter__": ("am_aiter", native_slot),
     "__anext__": ("am_anext", native_slot),
+    # Generator and coroutine classes get an am_send slot based on the generator helper
+    # method, so that PyIter_Send can drive them without raising StopIteration.
+    GENERATOR_HELPER_NAME: ("am_send", generate_am_send_wrapper),
 }
 
 SIDE_TABLES = [

mypyc/codegen/emitmodule.py

@@ -47,6 +47,7 @@
 from mypyc.codegen.literals import Literals
 from mypyc.common import (
     EXT_SUFFIX,
+    GENERATOR_HELPER_NAME,
     IS_FREE_THREADED,
     MODULE_PREFIX,
     PREFIX,
@@ -75,7 +76,7 @@
 from mypyc.ir.rtypes import RType
 from mypyc.irbuild.main import build_ir
 from mypyc.irbuild.mapper import Mapper
-from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME, load_type_map
+from mypyc.irbuild.prepare import load_type_map
 from mypyc.namegen import NameGenerator, exported_name
 from mypyc.options import CompilerOptions
 from mypyc.transform.copy_propagation import do_copy_propagation

mypyc/codegen/emitwrapper.py

@@ -595,6 +595,46 @@ def generate_len_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     return name
 
 
+def generate_am_send_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
+    """Generate an am_send slot for a generator or coroutine class.
+
+    This implements the C-level send protocol (used by PyIter_Send), which reports
+    normal completion without raising StopIteration.
+    """
+    name = f"{DUNDER_PREFIX}am_send{cl.name_prefix(emitter.names)}"
+    emitter.emit_line(
+        f"static PySendResult {name}(PyObject *self, PyObject *arg, PyObject **result) {{"
+    )
+    emitter.emit_line("PyObject *stop_iter_value = NULL;")
+    emitter.emit_line(
+        "PyObject *retval = {}{}{}(self, Py_None, Py_None, Py_None, arg, &stop_iter_value);".format(
+            emitter.get_group_prefix(fn.decl), NATIVE_PREFIX, fn.cname(emitter.names)
+        )
+    )
+    emitter.emit_line("if (retval != NULL) {")
+    emitter.emit_line("*result = retval;")
+    emitter.emit_line("return PYGEN_NEXT;")
+    emitter.emit_line("}")
+    emitter.emit_line("if (stop_iter_value == NULL) {")
+    # The helper raised an exception instead of using the out parameter. StopIteration
+    # signals normal completion when the generator or coroutine is already exhausted.
+    # TODO: Convert an explicit StopIteration raised in the body to RuntimeError instead
+    #       of treating it as normal completion.
+    emitter.emit_line("if (PyErr_ExceptionMatches(PyExc_StopIteration)) {")
+    emitter.emit_line("stop_iter_value = CPy_FetchStopIterationValue();")
+    emitter.emit_line("}")
+    emitter.emit_line("if (stop_iter_value == NULL) {")
+    emitter.emit_line("*result = NULL;")
+    emitter.emit_line("return PYGEN_ERROR;")
+    emitter.emit_line("}")
+    emitter.emit_line("}")
+    emitter.emit_line("*result = stop_iter_value;")
+    emitter.emit_line("return PYGEN_RETURN;")
+    emitter.emit_line("}")
+
+    return name
+
+
 def generate_bool_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     """Generates a wrapper for native __bool__ methods."""
     name = f"{DUNDER_PREFIX}{fn.name}{cl.name_prefix(emitter.names)}"

mypyc/common.py

@@ -20,6 +20,7 @@
 
 ENV_ATTR_NAME: Final = "__mypyc_env__"
 NEXT_LABEL_ATTR_NAME: Final = "__mypyc_next_label__"
+GENERATOR_HELPER_NAME: Final = "__mypyc_generator_helper__"
 TEMP_ATTR_NAME: Final = "__mypyc_temp__"
 LAMBDA_NAME: Final = "__mypyc_lambda__"
 PROPSET_PREFIX: Final = "__mypyc_setter__"

mypyc/irbuild/classdef.py

@@ -30,7 +30,7 @@
     is_class_var,
 )
 from mypy.types import Instance, UnboundType, get_proper_type
-from mypyc.common import MYPYC_DEFAULTS_SETUP, PROPSET_PREFIX
+from mypyc.common import GENERATOR_HELPER_NAME, MYPYC_DEFAULTS_SETUP, PROPSET_PREFIX
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
 from mypyc.ir.func_ir import FuncDecl, FuncSignature
 from mypyc.ir.ops import (
@@ -58,7 +58,6 @@
     handle_non_ext_method,
     load_type,
 )
-from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.util import (
     dataclass_type,
     default_attr_name,

mypyc/irbuild/for_helpers.py

@@ -28,7 +28,7 @@
     Var,
 )
 from mypy.types import LiteralType, TupleType, get_proper_type, get_proper_types
-from mypyc.common import IS_FREE_THREADED
+from mypyc.common import GENERATOR_HELPER_NAME, IS_FREE_THREADED
 from mypyc.ir.ops import (
     ERR_NEVER,
     BasicBlock,
@@ -68,7 +68,6 @@
 )
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.constant_fold import constant_fold_expr
-from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
 from mypyc.irbuild.vec import vec_append, vec_create, vec_get_item_unsafe, vec_init_item_unsafe
 from mypyc.primitives.dict_ops import (

mypyc/irbuild/generator.py

@@ -13,7 +13,12 @@
 from collections.abc import Callable
 
 from mypy.nodes import ARG_OPT, FuncDef, Var
-from mypyc.common import ENV_ATTR_NAME, GENERATOR_ATTRIBUTE_PREFIX, NEXT_LABEL_ATTR_NAME
+from mypyc.common import (
+    ENV_ATTR_NAME,
+    GENERATOR_ATTRIBUTE_PREFIX,
+    GENERATOR_HELPER_NAME,
+    NEXT_LABEL_ATTR_NAME,
+)
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FuncDecl, FuncIR
 from mypyc.ir.ops import (
@@ -51,7 +56,6 @@
     setup_func_for_recursive_call,
 )
 from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl, gen_generator_func_cleanup
-from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.targets import AssignmentTargetAttr
 from mypyc.primitives.exc_ops import (
     error_catch_op,

mypyc/irbuild/prepare.py

@@ -15,7 +15,7 @@
 
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import Final, NamedTuple
+from typing import NamedTuple
 
 from mypy.build import Graph
 from mypy.nodes import (
@@ -42,6 +42,7 @@
 from mypy.types import Instance, Type, get_proper_type
 from mypyc.common import (
     FAST_PREFIX,
+    GENERATOR_HELPER_NAME,
     MYPYC_DEFAULTS_SETUP,
     PROPSET_PREFIX,
     SELF_NAME,
@@ -83,8 +84,6 @@
 from mypyc.options import CompilerOptions
 from mypyc.sametype import is_same_type
 
-GENERATOR_HELPER_NAME: Final = "__mypyc_generator_helper__"
-
 
 def build_type_map(
     mapper: Mapper,

mypyc/irbuild/statement.py

@@ -49,7 +49,12 @@
     YieldExpr,
     YieldFromExpr,
 )
-from mypyc.common import KEEP_ALIVE_SHORT_LIVED, KEEP_ALIVE_WHOLE_EXPRESSION, TEMP_ATTR_NAME
+from mypyc.common import (
+    GENERATOR_HELPER_NAME,
+    KEEP_ALIVE_SHORT_LIVED,
+    KEEP_ALIVE_WHOLE_EXPRESSION,
+    TEMP_ATTR_NAME,
+)
 from mypyc.ir.ops import (
     ERR_NEVER,
     NAMESPACE_MODULE,
@@ -100,7 +105,6 @@
     FinallyNonlocalControl,
     TryFinallyNonlocalControl,
 )
-from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.specialize import apply_dunder_specialization
 from mypyc.irbuild.targets import (
     AssignmentTarget,
@@ -121,10 +125,9 @@
     reraise_exception_op,
     restore_exc_info_op,
 )
-from mypyc.primitives.generic_ops import iter_op, next_raw_op, py_delattr_op
+from mypyc.primitives.generic_ops import iter_op, py_delattr_op
 from mypyc.primitives.librt_threading_ops import lock_acquire_op, lock_release_op
 from mypyc.primitives.misc_ops import (
-    check_stop_op,
     coro_op,
     get_native_attrs_op,
     import_from_many_op,
@@ -1337,8 +1340,7 @@ def emit_yield_from_or_await(
         # This allows two optimizations:
         # 1) No need to call CPy_GetCoro() or iter() since for native generators
         #    it just returns the generator object (implemented here).
-        # 2) Instead of calling next(), call generator helper method directly,
-        #    since next() just calls __next__ which calls the helper method.
+        # 2) Call the generator helper method directly instead of using PyIter_Send.
         iter_val: Value = val
     else:
         get_op = coro_op if is_await else iter_op
@@ -1355,9 +1357,16 @@ def emit_yield_from_or_await(
         helper_method
     )
 
-    # Register where a native child stores its return value instead of raising
-    # StopIteration (only used on the fast path).
-    stop_iter_val = Register(object_rprimitive) if fast_path else None
+    # Neither the direct helper call nor CPyIter_Send raises StopIteration on normal
+    # completion. Instead they store the return value through a PyObject ** argument,
+    # which must be cleared before each call.
+    stop_iter_val = Register(object_rprimitive)
+
+    def clear_stop_iter_val() -> Value:
+        """Reset the stop iteration value and return a pointer to it."""
+        err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
+        builder.assign(stop_iter_val, err, line)
+        return builder.add(LoadAddress(object_pointer_rprimitive, stop_iter_val))
 
     def native_step(sent: Value) -> Value:
         """Advance a native generator/coroutine child by calling its helper directly.
@@ -1366,12 +1375,9 @@ def native_step(sent: Value) -> Value:
         raised. On normal completion the return value is stored in stop_iter_val,
         which is set to the error value if a real exception was raised instead.
         """
-        assert stop_iter_val is not None
         obj = builder.read(iter_reg, line)
         nn = builder.none_object()
-        err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
-        builder.assign(stop_iter_val, err, line)
-        ptr = builder.add(LoadAddress(object_pointer_rprimitive, stop_iter_val))
+        ptr = clear_stop_iter_val()
         m = MethodCall(obj, helper_method, [nn, nn, nn, sent, ptr], line)
         # Generators have custom error handling, so disable normal error handling.
         m.error_kind = ERR_NEVER
@@ -1384,19 +1390,19 @@ def native_step(sent: Value) -> Value:
         # an extra PyObject ** argument to helper where the stop iteration value is stored.
         _y_init = native_step(builder.none_object())
     else:
-        _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg, line)], line)
+        # Use PyIter_Send (via CPyIter_Send) so that a natively compiled iterator can
+        # complete through its am_send slot without raising StopIteration, even if we
+        # don't know the type statically.
+        ptr = clear_stop_iter_val()
+        _y_init = builder.call_c(
+            send_op, [builder.read(iter_reg, line), builder.none_object(), ptr], line
+        )
 
     builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR))
 
     builder.activate_block(stop_block)
-    if fast_path:
-        assert stop_iter_val is not None
-        builder.primitive_op(propagate_if_error_op, [stop_iter_val], line)
-        builder.assign(result, stop_iter_val, line)
-    else:
-        # Try extracting a return value from a StopIteration and return it.
-        # If it wasn't, this reraises the exception.
-        builder.assign(result, builder.call_c(check_stop_op, [], line), line)
+    builder.primitive_op(propagate_if_error_op, [stop_iter_val], line)
+    builder.assign(result, stop_iter_val, line)
     # Clear the spilled iterator/coroutine so that it will be freed.
     # Otherwise, the freeing of the spilled register would likely be delayed.
     err = builder.add(LoadErrorValue(iter_reg.type))
@@ -1446,10 +1452,13 @@ def else_body() -> None:
             # completion doesn't have to go through .send() and StopIteration.
             _y = native_step(builder.read(received_reg, line))
         else:
-            # Do a next() or a .send(). It will return NULL on exception
-            # but it won't automatically propagate.
+            # Do a next() or a .send(). It will return NULL on exception or on normal
+            # completion, but it won't automatically propagate an error.
+            ptr = clear_stop_iter_val()
             _y = builder.call_c(
-                send_op, [builder.read(iter_reg, line), builder.read(received_reg, line)], line
+                send_op,
+                [builder.read(iter_reg, line), builder.read(received_reg, line), ptr],
+                line,
             )
         ok, stop = BasicBlock(), BasicBlock()
         builder.add(Branch(_y, stop, ok, Branch.IS_ERROR))
@@ -1459,17 +1468,12 @@ def else_body() -> None:
         builder.assign(to_yield_reg, _y, line)
         builder.nonlocal_control[-1].gen_continue(builder, line)
 
+        # The iterator completed (or raised). If it completed normally, the
+        # return value is the result of the await/yield from. Otherwise
+        # propagate the exception.
         builder.activate_block(stop)
-        if fast_path:
-            assert stop_iter_val is not None
-            # The child either returned a value through the out pointer, or raised
-            # a real exception (in which case this propagates it).
-            builder.primitive_op(propagate_if_error_op, [stop_iter_val], line)
-            builder.assign(result, stop_iter_val, line)
-        else:
-            # Try extracting a return value from a StopIteration and return it.
-            # If it wasn't, this rereaises the exception.
-            builder.assign(result, builder.call_c(check_stop_op, [], line), line)
+        builder.primitive_op(propagate_if_error_op, [stop_iter_val], line)
+        builder.assign(result, stop_iter_val, line)
         builder.nonlocal_control[-1].gen_break(builder, line)
 
     builder.push_loop_stack(loop_block, done_block)

mypyc/lib-rt/CPy.h

@@ -992,7 +992,7 @@ static inline PyObject *CPy_TYPE(PyObject *obj) {
 PyObject *CPy_CalculateMetaclass(PyObject *type, PyObject *o);
 PyObject *CPy_GetCoro(PyObject *obj);
 PyObject *CPyGen_AlreadyExecutingError(int is_coroutine);
-PyObject *CPyIter_Send(PyObject *iter, PyObject *val);
+PyObject *CPyIter_Send(PyObject *iter, PyObject *val, PyObject **stop_iter_value);
 int CPy_YieldFromErrorHandle(PyObject *iter, PyObject **outp);
 PyObject *CPy_FetchStopIterationValue(void);
 PyObject *CPyType_FromTemplate(PyObject *template_,

mypyc/lib-rt/misc_ops.c

@@ -31,14 +31,24 @@ PyObject *CPyGen_AlreadyExecutingError(int is_coroutine)
     return NULL;
 }
 
-PyObject *CPyIter_Send(PyObject *iter, PyObject *val)
+PyObject *CPyIter_Send(PyObject *iter, PyObject *val, PyObject **stop_iter_value)
 {
     // Do a send, or a next if second arg is None.
     // (This behavior is to match the PEP 380 spec for yield from.)
-    if (Py_IsNone(val)) {
-        return CPyIter_Next(iter);
-    } else {
-        return PyObject_CallMethodOneArg(iter, mypyc_interned_str.send, val);
+    //
+    // If the iterator yields a value, return it. If it completes normally,
+    // return NULL and store the return value in *stop_iter_value, without
+    // raising StopIteration. If it raises, return NULL and leave
+    // *stop_iter_value untouched (the caller initializes it to NULL).
+    PyObject *result;
+    switch (PyIter_Send(iter, val, &result)) {
+        case PYGEN_NEXT:
+            return result;
+        case PYGEN_RETURN:
+            *stop_iter_value = result;
+            return NULL;
+        default:
+            return NULL;
     }
 }
 

mypyc/primitives/misc_ops.py

@@ -74,13 +74,15 @@
     error_kind=ERR_MAGIC,
 )
 
-# Do obj.send(value), or a next(obj) if second arg is None.
-# (This behavior is to match the PEP 380 spec for yield from.)
-# Like next_raw_op, don't swallow StopIteration,
-# but also don't propagate an error.
-# Can return NULL: see next_op.
+# Do obj.send(value), or a next(obj) if second arg is None, using the
+# PyIter_Send C API. (This behavior is to match the PEP 380 spec for yield from.)
+#
+# Returns the yielded value, or NULL if the iterator completed or raised. On
+# normal completion the return value is stored via the third (PyObject **)
+# argument instead of raising StopIteration. The caller must initialize the
+# pointed-to value to NULL, since it's only assigned on normal completion.
 send_op = custom_op(
-    arg_types=[object_rprimitive, object_rprimitive],
+    arg_types=[object_rprimitive, object_rprimitive, object_pointer_rprimitive],
     return_type=object_rprimitive,
     c_function_name="CPyIter_Send",
     error_kind=ERR_NEVER,