[mypyc] Avoid cyclic reference in nested functions (#16268)

Mypyc used to always put nested functions into the environment object,
which results in cyclic references, since the function object contains a 
reference to the environment.

Now we only do this if the body of a nested function refers to a nested
function (e.g. due to a recursive call). This means that in the majority 
of cases we can avoid the cyclic reference.

This speeds up self check by an impressive 7%. I'm not sure exactly why
the impact is so big, but spending less time in the cyclic garbage 
collector is probably a big part.

Author: JukkaL

mypyc/irbuild/builder.py

@@ -502,6 +502,11 @@ def non_function_scope(self) -> bool:
         # Currently the stack always has at least two items: dummy and top-level.
         return len(self.fn_infos) <= 2
 
+    def top_level_fn_info(self) -> FuncInfo | None:
+        if self.non_function_scope():
+            return None
+        return self.fn_infos[2]
+
     def init_final_static(
         self,
         lvalue: Lvalue,

mypyc/irbuild/context.py

@@ -22,6 +22,7 @@ def __init__(
         contains_nested: bool = False,
         is_decorated: bool = False,
         in_non_ext: bool = False,
+        add_nested_funcs_to_env: bool = False,
     ) -> None:
         self.fitem = fitem
         self.name = name
@@ -47,6 +48,7 @@ def __init__(
         self.contains_nested = contains_nested
         self.is_decorated = is_decorated
         self.in_non_ext = in_non_ext
+        self.add_nested_funcs_to_env = add_nested_funcs_to_env
 
         # TODO: add field for ret_type: RType = none_rprimitive
 

mypyc/irbuild/env_class.py

@@ -107,7 +107,7 @@ def load_env_registers(builder: IRBuilder) -> None:
         load_outer_envs(builder, fn_info.callable_class)
         # If this is a FuncDef, then make sure to load the FuncDef into its own environment
         # class so that the function can be called recursively.
-        if isinstance(fitem, FuncDef):
+        if isinstance(fitem, FuncDef) and fn_info.add_nested_funcs_to_env:
             setup_func_for_recursive_call(builder, fitem, fn_info.callable_class)
 
 

mypyc/irbuild/function.py

@@ -19,6 +19,7 @@
     ArgKind,
     ClassDef,
     Decorator,
+    FuncBase,
     FuncDef,
     FuncItem,
     LambdaExpr,
@@ -222,6 +223,7 @@ def c() -> None:
     is_decorated = fitem in builder.fdefs_to_decorators
     is_singledispatch = fitem in builder.singledispatch_impls
     in_non_ext = False
+    add_nested_funcs_to_env = has_nested_func_self_reference(builder, fitem)
     class_name = None
     if cdef:
         ir = builder.mapper.type_to_ir[cdef.info]
@@ -234,14 +236,15 @@ def c() -> None:
         func_name = name
     builder.enter(
         FuncInfo(
-            fitem,
-            func_name,
-            class_name,
-            gen_func_ns(builder),
-            is_nested,
-            contains_nested,
-            is_decorated,
-            in_non_ext,
+            fitem=fitem,
+            name=func_name,
+            class_name=class_name,
+            namespace=gen_func_ns(builder),
+            is_nested=is_nested,
+            contains_nested=contains_nested,
+            is_decorated=is_decorated,
+            in_non_ext=in_non_ext,
+            add_nested_funcs_to_env=add_nested_funcs_to_env,
         )
     )
 
@@ -267,7 +270,13 @@ def c() -> None:
         builder.enter(fn_info)
         setup_env_for_generator_class(builder)
         load_outer_envs(builder, builder.fn_info.generator_class)
-        if builder.fn_info.is_nested and isinstance(fitem, FuncDef):
+        top_level = builder.top_level_fn_info()
+        if (
+            builder.fn_info.is_nested
+            and isinstance(fitem, FuncDef)
+            and top_level
+            and top_level.add_nested_funcs_to_env
+        ):
             setup_func_for_recursive_call(builder, fitem, builder.fn_info.generator_class)
         create_switch_for_generator_class(builder)
         add_raise_exception_blocks_to_generator_class(builder, fitem.line)
@@ -344,6 +353,20 @@ def c() -> None:
     return func_ir, func_reg
 
 
+def has_nested_func_self_reference(builder: IRBuilder, fitem: FuncItem) -> bool:
+    """Does a nested function contain a self-reference in its body?
+
+    If a nested function only has references in the surrounding function,
+    we don't need to add it to the environment.
+    """
+    if any(isinstance(sym, FuncBase) for sym in builder.free_variables.get(fitem, set())):
+        return True
+    return any(
+        has_nested_func_self_reference(builder, nested)
+        for nested in builder.encapsulating_funcs.get(fitem, [])
+    )
+
+
 def gen_func_ir(
     builder: IRBuilder,
     args: list[Register],
@@ -768,7 +791,7 @@ def get_func_target(builder: IRBuilder, fdef: FuncDef) -> AssignmentTarget:
         # Get the target associated with the previously defined FuncDef.
         return builder.lookup(fdef.original_def)
 
-    if builder.fn_info.is_generator or builder.fn_info.contains_nested:
+    if builder.fn_info.is_generator or builder.fn_info.add_nested_funcs_to_env:
         return builder.lookup(fdef)
 
     return builder.add_local_reg(fdef, object_rprimitive)