[mypyc] Borrow generator frame state more aggressively

mypyc/codegen/emitfunc.py

@@ -454,9 +454,10 @@ def emit_load_attr_take_ref(
         Final attributes are never rebound (no setter), so there is no concurrent writer
         and no use-after-free window; an owned read uses the cheaper CPy_GetAttrRefFinal
         (a plain load + incref). Borrowed reads keep the plain load: they are only emitted
-        for attributes safe to borrow on free-threaded builds (Final and vec attrs -- see
-        transform_member_expr in irbuild), whose values live as long as their container.
-        The default (GIL) build always takes the plain-load path and increfs separately.
+        when safe on free-threaded builds, either for Final and vec attributes in
+        transform_member_expr in irbuild or for private generator-frame attributes proven
+        safe by borrow_generator_attrs. The default (GIL) build always takes the plain-load
+        path and increfs separately.
 
         Thread-confined attributes also use plain loads; see
         ClassIR.attrs_are_thread_confined.

mypyc/codegen/emitmodule.py

@@ -79,6 +79,7 @@
 from mypyc.irbuild.prepare import load_type_map
 from mypyc.namegen import NameGenerator, exported_name
 from mypyc.options import CompilerOptions
+from mypyc.transform.borrow_generator_attrs import borrow_generator_attrs
 from mypyc.transform.copy_propagation import do_copy_propagation
 from mypyc.transform.exceptions import insert_exception_handling
 from mypyc.transform.flag_elimination import do_flag_elimination
@@ -294,6 +295,8 @@ def compile_scc_to_ir(
                 insert_uninit_checks(fn, compiler_options.strict_traceback_checks)
                 # Insert exception handling.
                 insert_exception_handling(fn, compiler_options.strict_traceback_checks)
+                if fn in generator_spill_owners:
+                    borrow_generator_attrs(fn, generator_spill_owners[fn])
                 # Insert reference count handling.
                 insert_ref_count_opcodes(fn)
 

mypyc/irbuild/expression.py

@@ -315,6 +315,8 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
 
     is_final = builder.is_final_native_attr_ref(expr)
     scope = KEEP_ALIVE_SHORT_LIVED
+    # TODO: Support borrowing in subexpressions evaluated before suspension, such as
+    #       `await f(obj.final_attr)` and `yield f(obj.final_attr)`.
     if (
         is_final
         and builder.expression_depth > 1

mypyc/test-data/run-async.test

@@ -180,7 +180,7 @@ async def sleep(delay: float) -> None: ...
 [typing fixtures/typing-full.pyi]
 
 [case testRunAsyncAwaitInVariousPositions]
-from typing import cast, Any
+from typing import Any, Generator, cast
 
 import asyncio
 
@@ -264,6 +264,33 @@ async def construct(s: str) -> str:
 async def test_construct() -> None:
     assert await construct("foo") == "foo!foo?"
 
+class Box:
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+class Suspend:
+    def __await__(self) -> Generator[Any, Any, int]:
+        yield None
+        return 7
+
+def combine(value: Box, number: int) -> str:
+    return value.value + str(number)
+
+async def frame_arg_across_await(value: Box) -> str:
+    # Evaluation reads value before suspending, but combine() uses it after resume.
+    return combine(value, await Suspend())
+
+def drive(coro: Any) -> Any:
+    assert coro.send(None) is None
+    try:
+        coro.send(None)
+    except StopIteration as exc:
+        return exc.value
+    assert False
+
+def test_generator_attribute_borrow_does_not_cross_await() -> None:
+    assert drive(frame_arg_across_await(Box("value"))) == "value7"
+
 [file asyncio/__init__.pyi]
 async def sleep(t: float) -> None: ...
 

mypyc/test-data/run-generators.test

@@ -657,6 +657,22 @@ def test_near_yield() -> None:
         pass
     assert c.foo.flag == False
 
+class Box:
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+def select(first: Box, second: Box) -> str:
+    return first.value + second.value
+
+def walrus_gen(value: Box, replacement: Box) -> Iterator[str]:
+    # The first argument is read before the walrus rebinds the same frame field,
+    # but it is consumed by select() only after the rebind.
+    yield select(value, (value := replacement))
+
+def test_generator_attribute_borrow_killed_by_walrus() -> None:
+    it = walrus_gen(Box("old"), Box("new"))
+    assert next(it) == "oldnew"
+
 [case testGeneratorEarlyReturnWithBorrows]
 from typing import Iterator
 class Bar: