[mypyc] Use generator helper fast path on resume

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]