[mypyc] Move non-captured generator state to private frames

mypyc/codegen/emitfunc.py

@@ -13,12 +13,12 @@
     c_array_initializer,
 )
 from mypyc.common import (
-    GENERATOR_ATTRIBUTE_PREFIX,
     HAVE_IMMORTAL,
     IS_FREE_THREADED,
     NATIVE_PREFIX,
     REG_PREFIX,
     RUNNING_FIELD,
+    source_name_from_generator_attribute,
 )
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD, FuncDecl, FuncIR, all_values
@@ -485,6 +485,7 @@ def visit_get_attr(self, op: GetAttr) -> None:
         rtype = op.class_type
         cl = rtype.class_ir
         attr_rtype, decl_cl = cl.attr_details(op.attr)
+        source_attr_name = source_name_from_generator_attribute(op.attr, decl_cl.fullname)
         prefer_method = cl.is_trait and attr_rtype.error_overlap
         if cl.get_method(op.attr, prefer_method=prefer_method):
             # Properties are essentially methods, so use vtable access for them
@@ -528,17 +529,15 @@ def visit_get_attr(self, op: GetAttr) -> None:
                     ):
                         # Generate code for the following branch here to avoid
                         # redundant branches in the generated code.
-                        self.emit_attribute_error(branch, cl.name, op.attr)
+                        self.emit_attribute_error(branch, cl.name, source_attr_name)
                         self.emit_line("goto %s;" % self.label(branch.true))
                         merged_branch = branch
                         self.emitter.emit_line("}")
                 if not merged_branch:
                     exc_class = "PyExc_AttributeError"
                     self.emitter.emit_line(
                         'PyErr_SetString({}, "attribute {} of {} undefined");'.format(
-                            exc_class,
-                            repr(op.attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX)),
-                            repr(cl.name),
+                            exc_class, repr(source_attr_name), repr(cl.name)
                         )
                     )
 
@@ -1077,7 +1076,7 @@ def emit_traceback(self, op: Branch) -> None:
         if op.traceback_entry is not None:
             self.emitter.emit_traceback(self.source_path, self.module_name, op.traceback_entry)
 
-    def emit_attribute_error(self, op: Branch, class_name: str, attr: str) -> None:
+    def emit_attribute_error(self, op: Branch, class_name: str, source_attr_name: str) -> None:
         assert op.traceback_entry is not None
         if self.emitter.context.strict_traceback_checks:
             assert (
@@ -1090,7 +1089,7 @@ def emit_attribute_error(self, op: Branch, class_name: str, attr: str) -> None:
                 self.source_path.replace("\\", "\\\\"),
                 op.traceback_entry[0],
                 class_name,
-                attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX),
+                source_attr_name,
                 op.traceback_entry[1],
                 globals_static,
             )

mypyc/common.py

@@ -6,6 +6,7 @@
 from typing import Any, Final
 
 from mypy.util import unnamed_function
+from mypyc.namegen import exported_name
 
 PREFIX: Final = "CPyPy_"  # Python wrappers
 NATIVE_PREFIX: Final = "CPyDef_"  # Native functions etc.
@@ -27,8 +28,27 @@
 SELF_NAME: Final = "__mypyc_self__"
 MYPYC_DEFAULTS_SETUP: Final = "__mypyc_defaults_setup"
 GENERATOR_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_attribute__"
+GENERATOR_FRAME_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_frame_attribute__"
 CPYFUNCTION_NAME = "__cpyfunction__"
 
+
+def generator_frame_attribute_prefix(class_fullname: str, *, is_final_class: bool) -> str:
+    """Return the source-attribute prefix private to a generator frame class."""
+    if is_final_class:
+        return GENERATOR_FRAME_ATTRIBUTE_PREFIX
+    return f"{GENERATOR_FRAME_ATTRIBUTE_PREFIX}{exported_name(class_fullname)}_"
+
+
+def source_name_from_generator_attribute(name: str, class_fullname: str) -> str:
+    """Recover a source name from a generator frame or closure attribute."""
+    qualified_prefix = generator_frame_attribute_prefix(class_fullname, is_final_class=False)
+    if name.startswith(qualified_prefix):
+        return name.removeprefix(qualified_prefix)
+    if name.startswith(GENERATOR_FRAME_ATTRIBUTE_PREFIX):
+        return name.removeprefix(GENERATOR_FRAME_ATTRIBUTE_PREFIX)
+    return name.removeprefix(GENERATOR_ATTRIBUTE_PREFIX)
+
+
 # Omits the prefix added to user attribute fields, so it cannot collide with one.
 RUNNING_FIELD: Final = "mypyc_running"
 

mypyc/irbuild/builder.py

@@ -27,6 +27,7 @@
     DictionaryComprehension,
     Expression,
     FuncDef,
+    FuncItem,
     GeneratorExpr,
     IndexExpr,
     IntExpr,
@@ -75,6 +76,7 @@
     KEEP_ALIVE_WHOLE_EXPRESSION,
     MODULE_PREFIX,
     SELF_NAME,
+    generator_frame_attribute_prefix,
     shared_lib_name,
 )
 from mypyc.crash import catch_errors
@@ -763,20 +765,23 @@ def get_assignment_target(
                         reg_type = self.type_to_rtype(symbol.type)
                     else:
                         reg_type = self.node_type(lvalue)
-                    # A deleted error-overlap value needs the environment's
-                    # definedness bitmap. Other generator locals start in
-                    # registers and are promoted later if they cross a yield.
+                    # A deleted error-overlap value needs a definedness bitmap. Other generator
+                    # locals start in registers and are promoted later if they cross a yield.
                     if (
                         self.fn_info.is_generator
                         and reg_type.error_overlap
                         and symbol in self.deleted_vars
                     ):
-                        return self.add_var_to_env_class(
-                            symbol,
-                            reg_type,
-                            self.fn_info.generator_class,
-                            reassign=False,
-                            prefix=GENERATOR_ATTRIBUTE_PREFIX,
+                        if self.is_captured_by_nested_func(symbol):
+                            return self.add_var_to_env_class(
+                                symbol,
+                                reg_type,
+                                self.fn_info.generator_class,
+                                reassign=False,
+                                prefix=GENERATOR_ATTRIBUTE_PREFIX,
+                            )
+                        return self.add_var_to_generator_frame(
+                            symbol, reg_type, self.fn_info.generator_class.self_reg, reassign=False
                         )
 
                     return self.add_local_reg(symbol, reg_type)
@@ -1602,6 +1607,44 @@ def add_var_to_env_class(
             prefix=prefix,
         )
 
+    def is_free_variable_in_nested_func(self, fitem: FuncItem, symbol: SymbolNode) -> bool:
+        for nested in self.encapsulating_funcs.get(fitem, []):
+            if symbol in self.free_variables.get(nested, set()):
+                return True
+            if self.is_free_variable_in_nested_func(nested, symbol):
+                return True
+        return False
+
+    def is_captured_by_nested_func(self, symbol: SymbolNode) -> bool:
+        """Does a binding need to be visible to a nested function?"""
+        return symbol in self.free_variables.get(
+            self.fn_info.fitem, set()
+        ) or self.is_free_variable_in_nested_func(self.fn_info.fitem, symbol)
+
+    def add_var_to_generator_frame(
+        self,
+        var: SymbolNode,
+        rtype: RType,
+        frame_reg: Value,
+        reassign: bool = False,
+        always_defined: bool = False,
+        keep_alive_on_completion: bool = False,
+    ) -> AssignmentTarget:
+        """Add a generator-owned source binding to the private generator frame."""
+        cls = self.fn_info.generator_class.ir
+        return self.add_var_to_class(
+            var,
+            rtype,
+            cls,
+            frame_reg,
+            reassign=reassign,
+            always_defined=always_defined,
+            keep_alive_on_completion=keep_alive_on_completion,
+            prefix=generator_frame_attribute_prefix(
+                cls.fullname, is_final_class=cls.is_final_class
+            ),
+        )
+
     def add_var_to_class(
         self,
         var: SymbolNode,