diff --git a/mypyc/codegen/emitclass.py b/mypyc/codegen/emitclass.py index 8a1e3a53c86c9..9ecdea4327d95 100644 --- a/mypyc/codegen/emitclass.py +++ b/mypyc/codegen/emitclass.py @@ -252,7 +252,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None: getseters_name = f"{name_prefix}_getseters" vtable_name = f"{name_prefix}_vtable" traverse_name = f"{name_prefix}_traverse" - clear_name = f"{name_prefix}_clear" + clear_name = emitter.native_function_name(cl.clear) dealloc_name = f"{name_prefix}_dealloc" methods_name = f"{name_prefix}_methods" vtable_setup_name = f"{name_prefix}_trait_vtable_setup" @@ -280,7 +280,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None: fields["tp_dealloc"] = f"(destructor){name_prefix}_dealloc" if not cl.is_acyclic: fields["tp_traverse"] = f"(traverseproc){name_prefix}_traverse" - fields["tp_clear"] = f"(inquiry){name_prefix}_clear" + fields["tp_clear"] = f"(inquiry){clear_name}" # Populate .tp_finalize and generate a finalize method only if __del__ is defined for this class. del_method = next((e.method for e in cl.vtable_entries if e.name == "__del__"), None) if del_method: @@ -353,7 +353,11 @@ def emit_line() -> None: if not cl.is_acyclic: generate_traverse_for_class(cl, traverse_name, emitter) emit_line() - generate_clear_for_class(cl, clear_name, emitter) + generate_clear_for_class(cl, cl.clear, emitter) + emit_line() + generate_clear_for_class( + cl, cl.clear_on_completion, emitter, skip_attrs=cl.attrs_to_keep_alive_on_completion + ) emit_line() generate_dealloc_for_class(cl, dealloc_name, clear_name, bool(del_method), emitter) emit_line() @@ -902,12 +906,21 @@ def generate_traverse_for_class(cl: ClassIR, func_name: str, emitter: Emitter) - emitter.emit_line("}") -def generate_clear_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -> None: - emitter.emit_line("static int") - emitter.emit_line(f"{func_name}({cl.struct_name(emitter.names)} *self)") +def generate_clear_for_class( + cl: ClassIR, func_decl: FuncDecl, emitter: Emitter, skip_attrs: set[str] | None = None +) -> None: + if skip_attrs is None: + skip_attrs = set() + emitter.emit_line("static " + native_function_header(func_decl, emitter)) emitter.emit_line("{") + emitter.emit_line( + f"{cl.struct_name(emitter.names)} *self = " + f"({cl.struct_name(emitter.names)} *)cpy_r_self;" + ) for base in reversed(cl.base_mro): for attr, rtype in base.attributes.items(): + if attr in skip_attrs: + continue emitter.emit_gc_clear(f"self->{emitter.attr(attr)}", rtype) base_args = "(PyObject *)self" if cl.builtin_base: @@ -963,7 +976,7 @@ def generate_dealloc_for_class( if not cl.is_acyclic: emitter.emit_line("PyObject_GC_UnTrack(self);") if cl.builtin_base: - emitter.emit_line(f"{clear_func_name}(self);") + emitter.emit_line(f"{clear_func_name}((PyObject *)self);") # For native subclasses of builtins such as dict, the base deallocator # is responsible for tearing down base-owned storage and freeing memory. # Re-track self if base is GC-aware to match cpython's subtype_dealloc. @@ -978,7 +991,7 @@ def generate_dealloc_for_class( emit_reuse_dealloc(cl, emitter) # The trashcan is needed to handle deep recursive deallocations emitter.emit_line(f"CPy_TRASHCAN_BEGIN(self, {dealloc_func_name})") - emitter.emit_line(f"{clear_func_name}(self);") + emitter.emit_line(f"{clear_func_name}((PyObject *)self);") emitter.emit_line("Py_TYPE(self)->tp_free((PyObject *)self);") emitter.emit_line("CPy_TRASHCAN_END(self)") emitter.emit_line("done: ;") diff --git a/mypyc/ir/class_ir.py b/mypyc/ir/class_ir.py index 6ea1eb072999d..3b54331cb2f07 100644 --- a/mypyc/ir/class_ir.py +++ b/mypyc/ir/class_ir.py @@ -7,7 +7,7 @@ from mypyc.common import PROPSET_PREFIX, JsonDict from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg from mypyc.ir.ops import DeserMaps, Value -from mypyc.ir.rtypes import RInstance, RType, deserialize_type, object_rprimitive +from mypyc.ir.rtypes import RInstance, RType, c_int_rprimitive, deserialize_type, object_rprimitive from mypyc.namegen import NameGenerator, exported_name # Some notes on the vtable layout: Each concrete class has a vtable @@ -143,8 +143,25 @@ def __init__( module_name, FuncSignature([RuntimeArg("type", object_rprimitive)], RInstance(self)), ) + self.clear = FuncDecl( + name + "_clear", + None, + module_name, + FuncSignature([RuntimeArg("self", RInstance(self))], c_int_rprimitive), + internal=True, + ) + self.clear_on_completion = FuncDecl( + name + "_clear_on_completion", + None, + module_name, + FuncSignature([RuntimeArg("self", RInstance(self))], c_int_rprimitive), + internal=True, + ) # Attributes defined in the class (not inherited) self.attributes: dict[str, RType] = {} + # Attributes that must survive generator/coroutine completion because + # escaped nested functions may still read them as closure variables. + self.attrs_to_keep_alive_on_completion: set[str] = set() # Final attributes defined in the class (not inherited) self.final_attributes: set[str] = set() # Deletable attributes @@ -412,8 +429,11 @@ def serialize(self) -> JsonDict: "_serializable": self._serializable, "builtin_base": self.builtin_base, "ctor": self.ctor.serialize(), + "clear": self.clear.serialize(), + "clear_on_completion": self.clear_on_completion.serialize(), # We serialize dicts as lists to ensure order is preserved "attributes": [(k, t.serialize()) for k, t in self.attributes.items()], + "attrs_to_keep_alive_on_completion": sorted(self.attrs_to_keep_alive_on_completion), "final_attributes": sorted(self.final_attributes), # We try to serialize a name reference, but if the decl isn't in methods # then we can't be sure that will work so we serialize the whole decl. @@ -474,7 +494,10 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR: ir._serializable = data["_serializable"] ir.builtin_base = data["builtin_base"] ir.ctor = FuncDecl.deserialize(data["ctor"], ctx) + ir.clear = FuncDecl.deserialize(data["clear"], ctx) + ir.clear_on_completion = FuncDecl.deserialize(data["clear_on_completion"], ctx) ir.attributes = {k: deserialize_type(t, ctx) for k, t in data["attributes"]} + ir.attrs_to_keep_alive_on_completion = set(data["attrs_to_keep_alive_on_completion"]) ir.final_attributes = set(data["final_attributes"]) ir.method_decls = { k: ctx.functions[v].decl if isinstance(v, str) else FuncDecl.deserialize(v, ctx) diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index 0b598a00889f5..f4e3745836cf5 100644 --- a/mypyc/irbuild/builder.py +++ b/mypyc/irbuild/builder.py @@ -1568,12 +1568,15 @@ def add_var_to_env_class( base: FuncInfo | ImplicitClass, reassign: bool = False, always_defined: bool = False, + keep_alive_on_completion: bool = False, prefix: str = "", ) -> AssignmentTarget: # First, define the variable name as an attribute of the environment class, and then # construct a target for that attribute. name = prefix + remangle_redefinition_name(var.name) self.fn_info.env_class.attributes[name] = rtype + if keep_alive_on_completion: + self.fn_info.env_class.attrs_to_keep_alive_on_completion.add(name) if always_defined: self.fn_info.env_class.attrs_with_defaults.add(name) attr_target = AssignmentTargetAttr(base.curr_env_reg, name) diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index 3128543d4cd2c..c8c5f7fa68593 100644 --- a/mypyc/irbuild/env_class.py +++ b/mypyc/irbuild/env_class.py @@ -17,7 +17,7 @@ def g() -> int: from __future__ import annotations -from mypy.nodes import Argument, FuncDef, SymbolNode, Var +from mypy.nodes import Argument, FuncDef, FuncItem, SymbolNode, Var from mypyc.common import ( BITMAP_BITS, ENV_ATTR_NAME, @@ -60,6 +60,8 @@ class is generated, the function environment has not yet been # If the function is nested, its environment class must contain an environment # attribute pointing to its encapsulating functions' environment class. env_class.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class) + if builder.fn_info.contains_nested: + env_class.attrs_to_keep_alive_on_completion.add(ENV_ATTR_NAME) env_class.mro = [env_class] builder.fn_info.env_class = env_class builder.classes.append(env_class) @@ -229,7 +231,17 @@ def add_args_to_env( rtype = builder.type_to_rtype(arg.variable.type) assert base is not None, "base cannot be None for adding nonlocal args" builder.add_var_to_env_class( - arg.variable, rtype, base, reassign=reassign, prefix=prefix + arg.variable, + rtype, + base, + reassign=reassign, + keep_alive_on_completion=( + is_free_variable(builder, arg.variable) + or is_free_variable_in_nested_func( + builder, builder.fn_info.fitem, arg.variable + ) + ), + prefix=prefix, ) @@ -256,7 +268,12 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None: if isinstance(var, Var): rtype = builder.type_to_rtype(var.type) builder.add_var_to_env_class( - var, rtype, env_for_func, reassign=False, prefix=prefix + var, + rtype, + env_for_func, + reassign=False, + keep_alive_on_completion=True, + prefix=prefix, ) if builder.fn_info.fitem in builder.encapsulating_funcs: @@ -267,10 +284,16 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None: # the same name and signature across conditional blocks # will generate different callable classes, so the callable # class that gets instantiated must be generic. + nested_prefix = prefix if nested_fn.is_generator or nested_fn.is_coroutine: - prefix = GENERATOR_ATTRIBUTE_PREFIX + nested_prefix = GENERATOR_ATTRIBUTE_PREFIX builder.add_var_to_env_class( - nested_fn, object_rprimitive, env_for_func, reassign=False, prefix=prefix + nested_fn, + object_rprimitive, + env_for_func, + reassign=False, + keep_alive_on_completion=is_free_variable(builder, nested_fn), + prefix=nested_prefix, ) @@ -308,3 +331,14 @@ def setup_func_for_recursive_call( def is_free_variable(builder: IRBuilder, symbol: SymbolNode) -> bool: fitem = builder.fn_info.fitem return fitem in builder.free_variables and symbol in builder.free_variables[fitem] + + +def is_free_variable_in_nested_func( + builder: IRBuilder, fitem: FuncItem, symbol: SymbolNode +) -> bool: + for nested in builder.encapsulating_funcs.get(fitem, []): + if symbol in builder.free_variables.get(nested, set()): + return True + if is_free_variable_in_nested_func(builder, nested, symbol): + return True + return False diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 82fc74f9b1ce5..3e7f18c91c753 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -23,6 +23,7 @@ Call, Goto, Integer, + LoadErrorValue, MethodCall, RaiseStandardError, Register, @@ -49,7 +50,7 @@ load_outer_envs, setup_func_for_recursive_call, ) -from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl +from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl, gen_generator_func_cleanup from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME from mypyc.primitives.exc_ops import ( error_catch_op, @@ -107,6 +108,8 @@ class that implements the function (each function gets a separate class). setup_func_for_recursive_call( builder, fitem, builder.fn_info.generator_class, prefix=GENERATOR_ATTRIBUTE_PREFIX ) + cleanup_on_error = BasicBlock() + builder.builder.push_error_handler(cleanup_on_error) create_switch_for_generator_class(builder) add_raise_exception_blocks_to_generator_class(builder, fitem.line) @@ -114,6 +117,11 @@ class that implements the function (each function gets a separate class). builder.accept(fitem.body) builder.maybe_add_implicit_return() + builder.builder.pop_error_handler() + + builder.activate_block(cleanup_on_error) + gen_generator_func_cleanup(builder, fitem.line) + builder.add(Return(builder.add(LoadErrorValue(object_rprimitive)))) populate_switch_for_generator_class(builder) diff --git a/mypyc/irbuild/nonlocalcontrol.py b/mypyc/irbuild/nonlocalcontrol.py index e0a23769770d7..009a1d6a986bc 100644 --- a/mypyc/irbuild/nonlocalcontrol.py +++ b/mypyc/irbuild/nonlocalcontrol.py @@ -8,10 +8,13 @@ from abc import abstractmethod from typing import TYPE_CHECKING +from mypyc.ir.class_ir import ClassIR from mypyc.ir.ops import ( + ERR_NEVER, NO_TRACEBACK_LINE_NO, BasicBlock, Branch, + Call, Goto, Integer, Register, @@ -90,10 +93,7 @@ class GeneratorNonlocalControl(BaseNonlocalControl): """Default nonlocal control in a generator function outside statements.""" def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None: - # Assign an invalid next label number so that the next time - # __next__ is called, we jump to the case in which - # StopIteration is raised. - builder.assign(builder.fn_info.generator_class.next_label_target, Integer(-1), line) + gen_generator_func_cleanup(builder, line) # Raise a StopIteration containing a field for the value that # should be returned. Before doing so, create a new block @@ -132,6 +132,25 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None: builder.add(Return(Integer(0, object_rprimitive))) +def gen_class_clear_on_completion(builder: IRBuilder, obj: Value, cl: ClassIR, line: int) -> None: + call = Call(cl.clear_on_completion, [obj], line) + call.error_kind = ERR_NEVER + builder.add(call) + + +def gen_generator_func_cleanup(builder: IRBuilder, line: int) -> None: + """Clear references held by a completed generator or coroutine.""" + cls = builder.fn_info.generator_class + + # Assign an invalid next label number so that the next time __next__ is + # called, we jump to the case in which StopIteration is raised. + builder.assign(cls.next_label_target, Integer(-1), line) + + gen_class_clear_on_completion(builder, cls.curr_env_reg, builder.fn_info.env_class, line) + if builder.fn_info.env_class is not cls.ir: + gen_class_clear_on_completion(builder, cls.self_reg, cls.ir, line) + + class CleanupNonlocalControl(NonlocalControl): """Abstract nonlocal control that runs some cleanup code.""" diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index a7a08e7ced5f2..4c9fca54f4a2f 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -407,6 +407,12 @@ import asyncio import gc import platform +from testutil import is_gil_disabled + +def immediate_deallocation_checks_are_reliable() -> bool: + # Free-threaded builds can defer deallocation, so weakref liveness is not deterministic. + return not is_gil_disabled() + async def sleep() -> None: # Non-zero value hangs on Windows. time = 0 if platform.system() == "Windows" else 0.0001 @@ -512,6 +518,106 @@ async def stolen(n: int) -> int: async def test_stolen() -> None: await assert_no_leaks(lambda: stolen(200), 16) +async def retain_arg(c: set[int]) -> None: + async def nested() -> int: + await sleep() + return len(c) + + await nested() + +async def test_completed_coroutine_releases_environment() -> None: + if not immediate_deallocation_checks_are_reliable(): + return + + import weakref + + c = {1} + ref = weakref.ref(c) + coro = retain_arg(c) + c = {2} + await coro + assert ref() is None + +async def retain_arg_then_raise(c: set[int]) -> None: + async def nested() -> int: + await sleep() + return len(c) + + await nested() + raise ValueError + +async def test_failed_coroutine_releases_environment() -> None: + if not immediate_deallocation_checks_are_reliable(): + return + + import weakref + + c = {1} + ref = weakref.ref(c) + coro = retain_arg_then_raise(c) + c = {2} + try: + await coro + except ValueError: + pass + assert ref() is None + +async def close_retain_arg(c: set[int]) -> None: + await sleep() + len(c) + +async def test_closed_coroutine_releases_environment() -> None: + if not immediate_deallocation_checks_are_reliable(): + return + + import weakref + from typing import Any, cast + + c = {1} + ref = weakref.ref(c) + coro = cast(Any, close_retain_arg(c)) + coro.send(None) + c = {2} + coro.close() + assert ref() is None + +async def interpreted_leaf(c: set[int]) -> int: + await sleep() + return len(c) + +async def compiled_leaf(c: set[int]) -> int: + await sleep() + return len(c) + +async def compiled_awaits_interpreted(c: set[int]) -> int: + import interpreted + + return await interpreted.interpreted_leaf(c) + +async def interpreted_awaits_compiled(fn, c: set[int]) -> int: + return await fn(c) + +async def test_completed_mixed_coroutine_releases_environment() -> None: + if not immediate_deallocation_checks_are_reliable(): + return + + import interpreted + import weakref + + c = {1} + ref = weakref.ref(c) + coro = compiled_awaits_interpreted(c) + c = {2} + assert await coro == 1 + assert ref() is None + + c = {3} + ref = weakref.ref(c) + coro = interpreted.interpreted_awaits_compiled(compiled_leaf, c) + c = {4} + assert await coro == 1 + assert ref() is None + [file asyncio/__init__.pyi] async def sleep(t: float) -> None: ... @@ -1998,6 +2104,99 @@ for i in range(50): assert yields == ("hello",), yields assert val == "value" + str(i) + "world", repr(val) +[case testCoroutineClosureOutlivesCompletedCoroutine] +import asyncio +from typing import Callable + +async def make_reader() -> Callable[[], int]: + payload = [7] + await asyncio.sleep(0) + + def read_payload() -> int: + return payload[0] + + return read_payload + +async def make_arg_reader(payload: list[int]) -> Callable[[], int]: + await asyncio.sleep(0) + + def read_payload() -> int: + return payload[0] + + return read_payload + +async def make_indirect_reader() -> Callable[[], int]: + def leaf() -> int: + return 17 + + def reader() -> int: + return leaf() + + await asyncio.sleep(0) + return reader + +def test_closure_returned_from_completed_coroutine_keeps_capture() -> None: + read_payload = asyncio.run(make_reader()) + assert read_payload() == 7 + + payload = [13] + read_arg_payload = asyncio.run(make_arg_reader(payload)) + payload = [0] + assert read_arg_payload() == 13 + + indirect_reader = asyncio.run(make_indirect_reader()) + assert indirect_reader() == 17 + +async def test_closure_returned_from_awaited_coroutine_keeps_capture() -> None: + read_payload = await make_reader() + assert read_payload() == 7 + + payload = [13] + read_arg_payload = await make_arg_reader(payload) + payload = [0] + assert read_arg_payload() == 13 + + indirect_reader = await make_indirect_reader() + assert indirect_reader() == 17 + +[file asyncio/__init__.pyi] +from typing import Awaitable, TypeVar + +T = TypeVar("T") + +def run(x: Awaitable[T]) -> T: ... +async def sleep(t: float) -> None: ... + +[typing fixtures/typing-full.pyi] + +[case testNestedCoroutineOutlivesCompletedCoroutine] +import asyncio +from typing import Awaitable + +async def make_nested_coroutine() -> Awaitable[int]: + payload = [11] + await asyncio.sleep(0) + + async def read_payload() -> int: + await asyncio.sleep(0) + return payload[0] + + return read_payload() + +def test_nested_coroutine_returned_from_completed_coroutine_keeps_capture() -> None: + read_payload = asyncio.run(make_nested_coroutine()) + assert asyncio.run(read_payload) == 11 + +[file asyncio/__init__.pyi] +from typing import Awaitable, TypeVar + +T = TypeVar("T") + +def run(x: Awaitable[T]) -> T: ... +async def sleep(t: float) -> None: ... + +[typing fixtures/typing-full.pyi] + [case testBorrowedFinalAttrAcrossAsyncComprehension] import asyncio from typing import Final