[mypyc] Clear coroutine env on coroutine completion

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: ;")

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)

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)

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

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,13 +108,20 @@ 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)
 
     add_vars_to_env(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX)
 
     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)
 

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."""