[mypyc] Use generator helper fast path on resume (#21954)

Extend the direct helper method fast path for statically known native
generators and coroutines to resumes after suspension. This avoids
generic send dispatch and StopIteration handling on completion.

The savings compound in deep native coroutine chains, where each
await frame would otherwise pay this overhead.

In an (unrealistic) microbenchmark this made a long await chain
1.8x faster. With event loop overhead included and an
`asyncio.sleep(0)`, call, the improvement in modified 
microbenchmark was about 10% when using uvloop.

I used coding agent assist.

Author: JukkaL

mypyc/irbuild/statement.py

@@ -1351,30 +1351,46 @@ def emit_yield_from_or_await(
 
     stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock()
 
-    if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method):
-        # Second fast path optimization: call helper directly (see also comment above).
-        #
-        # Calling a generated generator, so avoid raising StopIteration by passing
-        # an extra PyObject ** argument to helper where the stop iteration value is stored.
-        fast_path = True
+    fast_path = isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(
+        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
+
+    def native_step(sent: Value) -> Value:
+        """Advance a native generator/coroutine child by calling its helper directly.
+
+        Returns the value yielded by the child, or NULL if the child completed or
+        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()
-        stop_iter_val = Register(object_rprimitive)
         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))
-        m = MethodCall(obj, helper_method, [nn, nn, nn, nn, ptr], line)
+        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
-        _y_init = builder.add(m)
+        return builder.add(m)
+
+    if fast_path:
+        # Second fast path optimization: call helper directly (see also comment above).
+        #
+        # Calling a generated generator, so avoid raising StopIteration by passing
+        # an extra PyObject ** argument to helper where the stop iteration value is stored.
+        _y_init = native_step(builder.none_object())
     else:
-        fast_path = False
         _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg, line)], 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:
@@ -1423,11 +1439,18 @@ def except_body() -> None:
         builder.nonlocal_control[-1].gen_break(builder, line)
 
     def else_body() -> None:
-        # Do a next() or a .send(). It will return NULL on exception
-        # but it won't automatically propagate.
-        _y = builder.call_c(
-            send_op, [builder.read(iter_reg, line), builder.read(received_reg, line)], line
-        )
+        # This path runs when the parent's yield is resumed normally via next() or send().
+        # An exception injected via throw() or close() takes the except_body path instead.
+        if fast_path:
+            # Reuse the direct helper call on resumes as well, so that native-to-native
+            # 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.
+            _y = builder.call_c(
+                send_op, [builder.read(iter_reg, line), builder.read(received_reg, line)], line
+            )
         ok, stop = BasicBlock(), BasicBlock()
         builder.add(Branch(_y, stop, ok, Branch.IS_ERROR))
 
@@ -1436,10 +1459,17 @@ def else_body() -> None:
         builder.assign(to_yield_reg, _y, line)
         builder.nonlocal_control[-1].gen_continue(builder, line)
 
-        # Try extracting a return value from a StopIteration and return it.
-        # If it wasn't, this rereaises the exception.
         builder.activate_block(stop)
-        builder.assign(result, builder.call_c(check_stop_op, [], line), line)
+        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.nonlocal_control[-1].gen_break(builder, line)
 
     builder.push_loop_stack(loop_block, done_block)

mypyc/test-data/run-async.test

@@ -2409,3 +2409,92 @@ async def test_borrow_final_attr_across_await_after_comprehension() -> None:
 async def sleep(t: float) -> None: ...
 
 [typing fixtures/typing-full.pyi]
+
+[case testRunAsyncResumedNativeAwait]
+from typing import Any, Generator
+
+from testutil import assertRaises
+
+class MyError(Exception):
+    pass
+
+class Suspend:
+    """Awaitable that suspends the given number of times before returning."""
+
+    def __init__(self, n: int) -> None:
+        self.n = n
+
+    def __await__(self) -> Generator[Any, Any, int]:
+        i = 0
+        while i < self.n:
+            yield i
+            i += 1
+        return self.n
+
+async def child(n: int) -> int:
+    total = await Suspend(n)
+    return total + 1
+
+async def parent(n: int) -> int:
+    # Statically known native child that suspends before completing, so the
+    # await loop must resume it after a suspension.
+    return await child(n) + 10
+
+async def grandparent(n: int) -> int:
+    return await parent(n) + 100
+
+async def tuple_child(n: int) -> tuple[int, str]:
+    await Suspend(n)
+    return (n, "x")
+
+async def tuple_parent(n: int) -> tuple[int, str]:
+    return await tuple_child(n)
+
+async def raising_child(n: int) -> int:
+    await Suspend(n)
+    raise MyError()
+
+async def raising_parent(n: int) -> int:
+    return await raising_child(n)
+
+async def catching_parent(n: int) -> int:
+    try:
+        return await raising_child(n)
+    except MyError:
+        return -1
+
+def drive(coro: Any) -> Any:
+    """Drive a coroutine to completion without an event loop."""
+    steps = 0
+    while True:
+        try:
+            coro.send(None)
+        except StopIteration as e:
+            return e.value
+        steps += 1
+        assert steps < 100
+
+def test_resumed_native_await() -> None:
+    for n in range(5):
+        assert drive(parent(n)) == n + 11
+        assert drive(grandparent(n)) == n + 111
+        assert drive(tuple_parent(n)) == (n, "x")
+
+def test_resumed_native_await_exception() -> None:
+    for n in range(5):
+        with assertRaises(MyError):
+            drive(raising_parent(n))
+        assert drive(catching_parent(n)) == -1
+
+def test_throw_into_suspended_native_child() -> None:
+    coro = parent(3)
+    coro.send(None)
+    with assertRaises(MyError):
+        coro.throw(MyError())
+
+def test_close_suspended_native_child() -> None:
+    coro = parent(3)
+    coro.send(None)
+    coro.close()
+
+[typing fixtures/typing-full.pyi]