diff --git a/mypyc/codegen/emitfunc.py b/mypyc/codegen/emitfunc.py index 0cef3fd78691d..e79faa41570ab 100644 --- a/mypyc/codegen/emitfunc.py +++ b/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. diff --git a/mypyc/codegen/emitmodule.py b/mypyc/codegen/emitmodule.py index b8039dcca823c..452b9b1c7c082 100644 --- a/mypyc/codegen/emitmodule.py +++ b/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) diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py index e3b9ad0a2b1a9..96c64092a95d0 100644 --- a/mypyc/irbuild/expression.py +++ b/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 diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 5eebb7348cf35..c6794e8b4b793 100644 --- a/mypyc/test-data/run-async.test +++ b/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: ... diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index 437b384c1fa43..e2ebf0a75b800 100644 --- a/mypyc/test-data/run-generators.test +++ b/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: diff --git a/mypyc/test/test_borrow_generator_attrs.py b/mypyc/test/test_borrow_generator_attrs.py new file mode 100644 index 0000000000000..d3c0878c1ec19 --- /dev/null +++ b/mypyc/test/test_borrow_generator_attrs.py @@ -0,0 +1,303 @@ +"""Tests for borrowing reads from private generator-frame attributes. + +These mostly use hand-built IR instead of normal data-driven tests, since generators +produce very verbose IR. +""" + +from __future__ import annotations + +import unittest + +from mypyc.common import GENERATOR_ATTRIBUTE_PREFIX, TEMP_ATTR_NAME +from mypyc.ir.class_ir import ClassIR +from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg +from mypyc.ir.ops import ( + BasicBlock, + Branch, + Call, + DecRef, + GetAttr, + Goto, + IncRef, + Integer, + KeepAlive, + LoadLiteral, + Op, + Register, + Return, + SetAttr, +) +from mypyc.ir.rtypes import ( + RInstance, + RTuple, + RType, + bit_rprimitive, + none_rprimitive, + str_rprimitive, +) +from mypyc.test.testutil import build_ir_for_single_file2 +from mypyc.transform.borrow_generator_attrs import borrow_generator_attrs +from mypyc.transform.exceptions import insert_exception_handling +from mypyc.transform.refcount import insert_ref_count_opcodes +from mypyc.transform.spill import insert_spills +from mypyc.transform.uninit import insert_uninit_checks + + +def return_none() -> Return: + return Return(Integer(1, none_rprimitive)) + + +class TestBorrowGeneratorAttrs(unittest.TestCase): + def make_helper( + self, blocks: list[BasicBlock], attributes: dict[str, RType] + ) -> tuple[FuncIR, ClassIR, Register]: + cl = ClassIR("gen", "module", is_generated=True) + cl.has_private_generator_frame = True + cl.has_running_flag = True + cl.attributes.update(attributes) + self_reg = Register(RInstance(cl), "self", is_arg=True) + decl = FuncDecl( + "__mypyc_generator_helper__", + cl.name, + cl.module_name, + FuncSignature([RuntimeArg("self", self_reg.type)], none_rprimitive), + ) + ir = FuncIR(decl, [self_reg], blocks) + cl.env_user_function = ir + return ir, cl, self_reg + + def make_single_block_helper( + self, body: list[Op], attributes: dict[str, RType] | None = None + ) -> tuple[FuncIR, ClassIR, Register]: + block = BasicBlock() + block.ops = [*body, return_none()] + return self.make_helper(blocks=[block], attributes=attributes or {"value": str_rprimitive}) + + def test_borrows_read_used_in_same_block(self) -> None: + ir, cl, self_reg = self.make_single_block_helper([]) + read = GetAttr(self_reg, "value", -1) + ir.blocks[0].ops[:0] = [read, KeepAlive([read])] + + borrow_generator_attrs(ir, cl) + + assert read.is_borrowed + + def test_requires_thread_confined_frame(self) -> None: + ir, cl, self_reg = self.make_single_block_helper([]) + read = GetAttr(self_reg, "value", -1) + ir.blocks[0].ops[:0] = [read, KeepAlive([read])] + cl.has_running_flag = False + + borrow_generator_attrs(ir, cl) + + assert not read.is_borrowed + + def test_same_attribute_store_kills_borrow(self) -> None: + ir, cl, self_reg = self.make_single_block_helper([]) + read = GetAttr(self_reg, "value", -1) + replacement = LoadLiteral("new", str_rprimitive) + ir.blocks[0].ops[:0] = [ + read, + replacement, + SetAttr(self_reg, "value", replacement, -1), + KeepAlive([read]), + ] + + borrow_generator_attrs(ir, cl) + + assert not read.is_borrowed + + def test_different_attribute_store_does_not_kill_borrow(self) -> None: + ir, cl, self_reg = self.make_single_block_helper( + [], {"value": str_rprimitive, "other": str_rprimitive} + ) + read = GetAttr(self_reg, "value", -1) + replacement = LoadLiteral("new", str_rprimitive) + ir.blocks[0].ops[:0] = [ + read, + replacement, + SetAttr(self_reg, "other", replacement, -1), + KeepAlive([read]), + ] + + borrow_generator_attrs(ir, cl) + + assert read.is_borrowed + + def test_arbitrary_call_does_not_kill_borrow(self) -> None: + ir, cl, self_reg = self.make_single_block_helper([]) + read = GetAttr(self_reg, "value", -1) + callee = FuncDecl("callee", None, "module", FuncSignature([], none_rprimitive)) + ir.blocks[0].ops[:0] = [read, Call(callee, [], -1), KeepAlive([read])] + + borrow_generator_attrs(ir, cl) + + assert read.is_borrowed + + def test_passing_frame_to_call_kills_borrow(self) -> None: + ir, cl, self_reg = self.make_single_block_helper([]) + read = GetAttr(self_reg, "value", -1) + ir.blocks[0].ops[:0] = [ + read, + Call(cl.clear_on_completion, [self_reg], -1), + KeepAlive([read]), + ] + + borrow_generator_attrs(ir, cl) + + assert not read.is_borrowed + + def test_storing_frame_in_attribute_kills_borrow(self) -> None: + ir, cl, self_reg = self.make_single_block_helper([]) + other_cl = ClassIR("other", "module") + other_cl.attributes["frame"] = self_reg.type + other = Register(RInstance(other_cl), "other") + read = GetAttr(self_reg, "value", -1) + ir.blocks[0].ops[:0] = [read, SetAttr(other, "frame", self_reg, -1), KeepAlive([read])] + + borrow_generator_attrs(ir, cl) + + assert not read.is_borrowed + + def test_borrows_across_basic_blocks(self) -> None: + first = BasicBlock() + second = BasicBlock() + ir, cl, self_reg = self.make_helper([first, second], {"value": str_rprimitive}) + read = GetAttr(self_reg, "value", -1) + first.ops = [read, Goto(second)] + second.ops = [KeepAlive([read]), return_none()] + + borrow_generator_attrs(ir, cl) + + assert read.is_borrowed + + def test_store_on_one_branch_kills_borrow_at_join(self) -> None: + entry = BasicBlock() + store = BasicBlock() + no_store = BasicBlock() + join = BasicBlock() + ir, cl, self_reg = self.make_helper( + [entry, store, no_store, join], {"value": str_rprimitive} + ) + read = GetAttr(self_reg, "value", -1) + replacement = LoadLiteral("new", str_rprimitive) + entry.ops = [read, Branch(Integer(1, bit_rprimitive), store, no_store, Branch.BOOL)] + store.ops = [replacement, SetAttr(self_reg, "value", replacement, -1), Goto(join)] + no_store.ops = [Goto(join)] + join.ops = [KeepAlive([read]), return_none()] + + borrow_generator_attrs(ir, cl) + + assert not read.is_borrowed + + def test_suspension_kills_borrow(self) -> None: + entry = BasicBlock() + before_yield = BasicBlock() + continuation = BasicBlock() + ir, cl, self_reg = self.make_helper( + [entry, before_yield, continuation], {"value": str_rprimitive} + ) + read = GetAttr(self_reg, "value", -1) + entry.ops = [Branch(Integer(1, bit_rprimitive), before_yield, continuation, Branch.BOOL)] + before_yield.ops = [read, Return(Integer(1, none_rprimitive), yield_target=continuation)] + continuation.ops = [KeepAlive([read]), return_none()] + + borrow_generator_attrs(ir, cl) + + assert not read.is_borrowed + read.is_borrowed = True + with self.assertRaisesRegex(AssertionError, "cannot spill borrowed attribute read"): + insert_spills(ir, cl) + + def test_borrows_nullable_rtuple_attribute(self) -> None: + tuple_type = RTuple([str_rprimitive, str_rprimitive]) + ir, cl, self_reg = self.make_single_block_helper([], {"value": tuple_type}) + read = GetAttr(self_reg, "value", -1, allow_error_value=True) + ir.blocks[0].ops[:0] = [read, KeepAlive([read])] + + borrow_generator_attrs(ir, cl) + + assert read.is_borrowed + + def test_skips_read_that_all_consumers_steal(self) -> None: + ir, cl, self_reg = self.make_single_block_helper( + [], {"value": str_rprimitive, "other": str_rprimitive} + ) + read = GetAttr(self_reg, "value", -1) + ir.blocks[0].ops[:0] = [read, SetAttr(self_reg, "other", read, -1)] + + borrow_generator_attrs(ir, cl) + + assert not read.is_borrowed + + def test_borrows_private_frame_but_not_separate_environment(self) -> None: + source = """\ +from typing import Generator + +def use(value: str) -> str: + return value + +def gen(value: str) -> Generator[str, None, str]: + other = value + def nested() -> str: + return value + yield use(other) + return use(other) +""" + module, _, _, _ = build_ir_for_single_file2(source.splitlines()) + cl = next(cl for cl in module.classes if cl.name == "gen_gen") + ir = cl.env_user_function + assert ir is not None + insert_uninit_checks(ir, True) + insert_exception_handling(ir, True) + + borrow_generator_attrs(ir, cl) + + reads = [op for block in ir.blocks for op in block.ops if isinstance(op, GetAttr)] + private_reads = [ + op + for op in reads + if op.obj is ir.arg_regs[0] and op.attr == GENERATOR_ATTRIBUTE_PREFIX + "other" + ] + shared_reads = [ + op + for op in reads + if op.obj is not ir.arg_regs[0] and op.attr == GENERATOR_ATTRIBUTE_PREFIX + "value" + ] + assert private_reads and all(op.is_borrowed for op in private_reads) + assert shared_reads and not any(op.is_borrowed for op in shared_reads) + + def test_borrows_promoted_saved_exception_state(self) -> None: + source = """\ +from typing import Iterator + +def gen() -> Iterator[int]: + try: + yield 1 + finally: + yield 2 +""" + module, _, _, _ = build_ir_for_single_file2(source.splitlines()) + cl = next(cl for cl in module.classes if cl.name == "gen_gen") + ir = cl.env_user_function + assert ir is not None + slots = {name for name in cl.attributes if name.startswith(TEMP_ATTR_NAME + "3_")} + insert_uninit_checks(ir, True) + insert_exception_handling(ir, True) + + borrow_generator_attrs(ir, cl) + + reads = [ + op + for block in ir.blocks + for op in block.ops + if isinstance(op, GetAttr) and op.attr in slots + ] + assert reads and all(op.allow_error_value and op.is_borrowed for op in reads) + insert_ref_count_opcodes(ir) + assert not any( + isinstance(op, IncRef | DecRef) and op.src in reads + for block in ir.blocks + for op in block.ops + ) + insert_spills(ir, cl) diff --git a/mypyc/transform/borrow_generator_attrs.py b/mypyc/transform/borrow_generator_attrs.py new file mode 100644 index 0000000000000..8ad8793130ce3 --- /dev/null +++ b/mypyc/transform/borrow_generator_attrs.py @@ -0,0 +1,156 @@ +"""Borrow reads from private generator-frame attributes. + +The frame is private, and the generator running flag prevents another helper invocation from +accessing it concurrently, so attribute values can be borrowed while the helper runs. +""" + +from __future__ import annotations + +from collections.abc import Set as AbstractSet +from typing import Final + +from mypyc.analysis.dataflow import MUST_ANALYSIS, BaseAnalysisVisitor, get_cfg, run_analysis +from mypyc.ir.class_ir import ClassIR +from mypyc.ir.func_ir import FuncIR +from mypyc.ir.ops import ( + Assign, + AssignMulti, + BasicBlock, + Branch, + GetAttr, + KeepAlive, + Op, + RegisterOp, + Return, + SetAttr, + SetMem, + Unreachable, + Value, +) +from mypyc.ir.rtypes import RInstance + +GenAndKill = tuple[AbstractSet[GetAttr], AbstractSet[GetAttr]] +EMPTY: GenAndKill = (frozenset(), frozenset()) + + +class BorrowGeneratorAttrsVisitor(BaseAnalysisVisitor[GetAttr]): + def __init__( + self, + self_reg: Value, + candidates: set[GetAttr], + candidates_by_attr: dict[str, set[GetAttr]], + ) -> None: + self.self_reg: Final = self_reg + self.candidates: Final = candidates + self.candidates_by_attr: Final = candidates_by_attr + + def visit_branch(self, op: Branch) -> GenAndKill: + return self.visit_sources(op) + + def visit_return(self, op: Return) -> GenAndKill: + if op.yield_target is not None: + return frozenset(), self.candidates + return self.visit_sources(op) + + def visit_unreachable(self, op: Unreachable) -> GenAndKill: + return EMPTY + + def visit_register_op(self, op: RegisterOp) -> GenAndKill: + return self.visit_sources(op) + + def visit_assign(self, op: Assign) -> GenAndKill: + return self.visit_sources(op) + + def visit_assign_multi(self, op: AssignMulti) -> GenAndKill: + return self.visit_sources(op) + + def visit_set_mem(self, op: SetMem) -> GenAndKill: + return self.visit_sources(op) + + def visit_get_attr(self, op: GetAttr) -> GenAndKill: + if op in self.candidates: + return {op}, frozenset() + return EMPTY + + def visit_set_attr(self, op: SetAttr) -> GenAndKill: + if op.src is self.self_reg: + return frozenset(), self.candidates + if op.obj is self.self_reg: + return frozenset(), self.candidates_by_attr.get(op.attr, frozenset()) + return EMPTY + + def visit_keep_alive(self, op: KeepAlive) -> GenAndKill: + return EMPTY + + def visit_sources(self, op: Op) -> GenAndKill: + # An operation that uses the frame may access its attributes in a way this analysis + # doesn't track, so conservatively invalidate all outstanding borrows. Safe direct frame + # operations are handled more precisely by the methods above. Generator helper IR is + # assumed not to let an operation retain the frame and access it after the operation. + if any(source is self.self_reg for source in op.sources()): + return frozenset(), self.candidates + return EMPTY + + +def borrow_generator_attrs(ir: FuncIR, cl: ClassIR) -> None: + """Mark safe reads from a private generator frame as borrowed. + + Only private frames without Python-visible attribute getsetters qualify, and the + generator running flag prevents concurrent helper invocations. A read can stay borrowed + until the helper writes the same attribute, passes the frame to an operation that may + access it, or suspends. The analysis assumes that no operation retains the frame; a frame + use invalidates existing borrows but does not prevent new borrows after the use. + """ + if not cl.attrs_are_thread_confined() or cl.env_user_function is not ir or not ir.arg_regs: + return + + self_reg = ir.arg_regs[0] + if not isinstance(self_reg.type, RInstance) or self_reg.type.class_ir is not cl: + return + + candidates = { + op + for block in ir.blocks + for op in block.ops + if isinstance(op, GetAttr) + and op.obj is self_reg + and op.attr in cl.attributes + and cl.attr_type(op.attr).is_refcounted + and not op.is_borrowed + } + if not candidates: + return + + use_positions: dict[GetAttr, list[tuple[BasicBlock, int]]] = { + candidate: [] for candidate in candidates + } + only_stolen = set(candidates) + for block in ir.blocks: + for index, op in enumerate(block.ops): + stolen = op.stolen() + for source in op.unique_sources(): + if source in candidates: + use_positions[source].append((block, index)) + if source not in stolen: + only_stolen.discard(source) + + candidates -= only_stolen + if not candidates: + return + + candidates_by_attr: dict[str, set[GetAttr]] = {} + for candidate in candidates: + candidates_by_attr.setdefault(candidate.attr, set()).add(candidate) + + result = run_analysis( + blocks=ir.blocks, + cfg=get_cfg(ir.blocks), + gen_and_kill=BorrowGeneratorAttrsVisitor(self_reg, candidates, candidates_by_attr), + initial=set(), + kind=MUST_ANALYSIS, + backward=False, + universe=candidates, + ) + for candidate in candidates: + if all(candidate in result.before[position] for position in use_positions[candidate]): + candidate.is_borrowed = True diff --git a/mypyc/transform/spill.py b/mypyc/transform/spill.py index 313e795262131..2ce701f450031 100644 --- a/mypyc/transform/spill.py +++ b/mypyc/transform/spill.py @@ -62,6 +62,12 @@ def spill_regs( spill_locs = {} # Sort values to make the order deterministic. for i, val in enumerate(sort_values(to_spill, blocks)): + # Borrowed literals can be spilled because their backing static reference + # remains alive. Attribute reads instead borrow from the frame and must not + # outlive this helper invocation. + assert not ( + isinstance(val, GetAttr) and val.is_borrowed + ), f"cannot spill borrowed attribute read {val}" # Generator classes for overriding methods can inherit from one another. Include the # module-qualified owning class name so unrelated helper spills don't alias an inherited # struct field.