[mypyc] Keep all compiler spills on generator frames (#21958)

They are internal to the generator, and by moving them to the generator
from the environment they can be accessed using faster attribute ops on
free-threaded builds.

This made a microbenchmark marginally faster on Python 3.14
(free-threaded).

I used coding agent assist.

Author: JukkaL

mypyc/irbuild/builder.py

@@ -303,7 +303,7 @@ def __init__(
         # Whether the current top-level expression contains a suspension point
         # (await, yield or yield from). A whole-expression borrow can't span such a
         # point, since the borrowed value (and its root) live in registers that are
-        # not spilled into the generator environment across the suspend.
+        # not spilled into the generator frame across the suspend.
         self.expr_has_suspend = False
         # Saved expression state for enclosing functions (see enter()/leave()).
         self.expression_depth_stack: list[int] = []
@@ -1046,22 +1046,26 @@ def pop_loop_stack(self) -> None:
         self.nonlocal_control.pop()
 
     def make_spill_target(self, type: RType) -> AssignmentTarget:
-        """Moves a given Value instance into the generator class' environment class."""
-        name = f"{TEMP_ATTR_NAME}{self.temp_counter}"
+        """Moves a given Value instance into the private generator frame."""
+        frame = self.fn_info.generator_class
+        # 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.
+        name = f"{TEMP_ATTR_NAME}1_{exported_name(frame.ir.fullname)}_{self.temp_counter}"
         self.temp_counter += 1
-        target = self.add_var_to_env_class(Var(name), type, self.fn_info.generator_class)
+        target = self.add_var_to_class(Var(name), type, frame.ir, frame.self_reg)
         return target
 
     def spill(self, value: Value) -> AssignmentTarget:
-        """Moves a given Value instance into the generator class' environment class."""
+        """Moves a given Value instance into the private generator frame."""
         target = self.make_spill_target(value.type)
         # Shouldn't be able to fail
         self.assign(target, value, NO_TRACEBACK_LINE_NO)
         return target
 
     def maybe_spill(self, value: Value) -> Value | AssignmentTarget:
         """
-        Moves a given Value instance into the environment class for generator functions. For
+        Moves a given Value instance into the private frame for generator functions. For
         non-generator functions, leaves the Value instance as it is.
 
         Returns an AssignmentTarget associated with the Value for generator functions and the
@@ -1073,7 +1077,7 @@ def maybe_spill(self, value: Value) -> Value | AssignmentTarget:
 
     def maybe_spill_assignable(self, value: Value) -> Register | AssignmentTarget:
         """
-        Moves a given Value instance into the environment class for generator functions. For
+        Moves a given Value instance into the private frame for generator functions. For
         non-generator functions, allocate a temporary Register.
 
         Returns an AssignmentTarget associated with the Value for generator functions and an
@@ -1633,24 +1637,45 @@ def add_var_to_env_class(
         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.
+        return self.add_var_to_class(
+            var,
+            rtype,
+            self.fn_info.env_class,
+            base.curr_env_reg,
+            reassign=reassign,
+            always_defined=always_defined,
+            keep_alive_on_completion=keep_alive_on_completion,
+            prefix=prefix,
+        )
+
+    def add_var_to_class(
+        self,
+        var: SymbolNode,
+        rtype: RType,
+        cls: ClassIR,
+        base: Value,
+        reassign: bool = False,
+        always_defined: bool = False,
+        keep_alive_on_completion: bool = False,
+        prefix: str = "",
+    ) -> AssignmentTarget:
+        """Declare an attribute on a class and construct a target using an explicit base."""
         name = prefix + remangle_redefinition_name(var.name)
-        self.fn_info.env_class.attributes[name] = rtype
+        cls.attributes[name] = rtype
         if keep_alive_on_completion:
-            self.fn_info.env_class.attrs_to_keep_alive_on_completion.add(name)
+            cls.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)
+            cls.attrs_with_defaults.add(name)
+        attr_target = AssignmentTargetAttr(base, name)
 
         if reassign:
             # Read the local definition of the variable, and set the corresponding attribute of
-            # the environment class' variable to be that value.
+            # the class' variable to be that value.
             reg = self.read(self.lookup(var), self.fn_info.fitem.line)
-            self.add(SetAttr(base.curr_env_reg, name, reg, self.fn_info.fitem.line))
+            self.add(SetAttr(base, name, reg, self.fn_info.fitem.line))
 
         # Override the local definition of the variable to instead point at the variable in
-        # the environment class.
+        # the class.
         return self.add_target(var, attr_target)
 
     def is_builtin_ref_expr(self, expr: RefExpr) -> bool:

mypyc/irbuild/for_helpers.py

@@ -705,7 +705,7 @@ def need_cleanup(self) -> bool:
     def init(self, expr_reg: Value, target_type: RType) -> None:
         # Define targets to contain the expression, along with the iterator that will be used
         # for the for-loop. If we are inside of a generator function, spill these into the
-        # environment class.
+        # private generator frame.
         builder = self.builder
         iter_reg = builder.primitive_op(iter_op, [expr_reg], self.line)
         builder.maybe_spill(expr_reg)
@@ -753,7 +753,7 @@ def need_cleanup(self) -> bool:
 
     def init(self, expr_reg: Value, target_type: RType) -> None:
         # Define target to contains the generator expression. It's also the iterator.
-        # If we are inside a generator function, spill these into the environment class.
+        # If we are inside a generator function, spill these into the private generator frame.
         builder = self.builder
         self.iter_target = builder.maybe_spill(expr_reg)
         self.target_type = target_type
@@ -811,7 +811,7 @@ def init(self, expr_reg: Value, target_type: RType) -> None:
         # Define targets to contain the expression, along with the
         # iterator that will be used for the for-loop. We are inside
         # of a generator function, so we will spill these into
-        # environment class.
+        # the private generator frame.
         builder = self.builder
         iter_reg = builder.call_c(aiter_op, [expr_reg], self.line)
         builder.maybe_spill(expr_reg)
@@ -910,7 +910,7 @@ def init(
         self.reverse = reverse
         # Define target to contain the expression, along with the index that will be used
         # for the for-loop. If we are inside of a generator function, spill these into the
-        # environment class.
+        # private generator frame.
         self.expr_target = builder.maybe_spill(expr_reg)
         if is_immutable_rprimitive(expr_reg.type):
             # If the expression is an immutable type, we can load the length just once.
@@ -1011,7 +1011,7 @@ def init(self, expr_reg: Value, target_type: RType) -> None:
         builder = self.builder
         self.target_type = target_type
 
-        # We add some variables to environment class, so they can be read across yield.
+        # Spill some values so they can be read across yield.
         self.expr_target = builder.maybe_spill(expr_reg)
         offset = Integer(0)
         self.offset_target = builder.maybe_spill_assignable(offset)

mypyc/test-data/run-generators.test

@@ -772,9 +772,12 @@ def make() -> str:
 
 def outer() -> Generator[str, str, str]:
     def nested(value: str) -> Generator[str, str, str]:
-        # The result of make() is a compiler-generated temporary live across
-        # the yield. The source argument stays in the separate environment.
-        return make() + (yield value)
+        for item in [value]:
+            # The loop iterator and the result of make() are compiler-generated
+            # temporaries live across the yield. The source bindings stay in the
+            # separate environment.
+            return make() + (yield item)
+        return "unreachable"
 
     return nested("right")
 
@@ -792,9 +795,12 @@ def make_number() -> int:
 
 class Same(other_base.Same):
     def gen(self) -> Generator[None, None, object]:
-        # This spill has a different native representation from the base class's spill.
+        # The builder and transform spills have different native representations from the
+        # corresponding base class spills.
         # Both generator frames have the same short name, but different full names.
-        return (make_number(), (yield None))[0]
+        for _ in range(1):
+            return (make_number(), (yield None))[0]
+        return None
 
 def run(g: Generator[None, None, object]) -> object:
     assert next(g) is None
@@ -816,8 +822,10 @@ def make_text() -> str:
 
 class Same:
     def gen(self) -> Generator[None, None, object]:
-        # The call result is live across the yield.
-        return (make_text(), (yield None))[0]
+        # The loop state and call result are live across the yield.
+        for _ in [None]:
+            return (make_text(), (yield None))[0]
+        return None
 
 [case testGeneratorReuse]
 from typing import Iterator, Any

mypyc/test/test_spill.py

@@ -24,7 +24,11 @@ def make() -> str:
 
 def outer():
     def nested(value: str):
-        return make() + (yield value)
+        for item in [value]:
+            # The loop iterator uses IR-builder-managed spill slots, while the
+            # result of make() is spilled later by the spill transform.
+            return make() + (yield item)
+        return "unreachable"
     return nested("right")
 """
         module, _, _, _ = build_ir_for_single_file2(source.splitlines())
@@ -41,6 +45,7 @@ def nested(value: str):
         # is protected by the running flag and can use plain attribute access.
         assert frame.attrs_are_thread_confined()
         assert NEXT_LABEL_ATTR_NAME in frame.attributes
+        assert any(name.startswith(TEMP_ATTR_NAME + "1_") for name in frame.attributes)
         assert any(name.startswith(TEMP_ATTR_NAME + "2_") for name in frame.attributes)
 
         # Source-level variables stay in the shared environment.