[mypyc] Move private generator spills from environ to generator objects

mypyc/codegen/emitmodule.py

@@ -279,11 +279,11 @@ def compile_scc_to_ir(
     if errors.num_errors > 0:
         return modules
 
-    env_user_functions = {}
+    generator_spill_owners = {}
     for module in modules.values():
         for cls in module.classes:
             if cls.env_user_function:
-                env_user_functions[cls.env_user_function] = cls
+                generator_spill_owners[cls.env_user_function] = cls
 
     for module in modules.values():
         module_path = result.graph[module.fullname].xpath
@@ -296,8 +296,8 @@ def compile_scc_to_ir(
                 # Insert reference count handling.
                 insert_ref_count_opcodes(fn)
 
-                if fn in env_user_functions:
-                    insert_spills(fn, env_user_functions[fn])
+                if fn in generator_spill_owners:
+                    insert_spills(fn, generator_spill_owners[fn])
 
                 if compiler_options.log_trace:
                     insert_event_trace_logging(fn, compiler_options)

mypyc/ir/class_ir.py

@@ -233,7 +233,8 @@ def __init__(
         # value of an attribute is the same as the error value.
         self.bitmap_attrs: list[str] = []
 
-        # If this is a generator environment class, what is the actual method for it
+        # If this class owns a generator helper's compiler-generated spill slots, what is the
+        # actual helper method for it.
         self.env_user_function: FuncIR | None = None
 
         # If True, keep one freed, cleared instance available for immediate reuse to
@@ -256,8 +257,10 @@ def __init__(
         # Does this generator or coroutine helper serialize execution using an instance flag?
         self.has_running_flag = False
 
-        # Does this generator object contain its merged environment?
-        self.has_merged_generator_env = False
+        # Are this generator object's implementation attributes only accessed while its
+        # running flag is held (or before publication/during destruction)? This also applies
+        # when captured source variables live in a separate environment object.
+        self.has_private_generator_frame = False
 
     def __repr__(self) -> str:
         return (
@@ -319,13 +322,13 @@ def needs_getseters_table(self) -> bool:
     def attrs_are_thread_confined(self) -> bool:
         """Can these attributes safely use plain access in free-threaded builds?
 
-        This requires locals to live directly in the generator object, execution to be
-        serialized by its running flag, and no Python getseters exposing the attributes.
-        A separate environment does not qualify because captured locals may be accessed
-        by nested functions.
+        This requires the attributes to belong to a private generator frame, execution to
+        be serialized by its running flag, and no Python getseters exposing the attributes.
+        Captured locals in a separate environment don't qualify, since nested functions may
+        access them independently.
         """
         return (
-            self.has_merged_generator_env
+            self.has_private_generator_frame
             and self.has_running_flag
             and not self.needs_getseters_table
         )
@@ -521,7 +524,7 @@ def serialize(self) -> JsonDict:
             "env_user_function": self.env_user_function.id if self.env_user_function else None,
             "reuse_freed_instance": self.reuse_freed_instance,
             "has_running_flag": self.has_running_flag,
-            "has_merged_generator_env": self.has_merged_generator_env,
+            "has_private_generator_frame": self.has_private_generator_frame,
             "is_acyclic": self.is_acyclic,
             "is_enum": self.is_enum,
             "is_coroutine": self.coroutine_name,
@@ -589,7 +592,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
         )
         ir.reuse_freed_instance = data["reuse_freed_instance"]
         ir.has_running_flag = data["has_running_flag"]
-        ir.has_merged_generator_env = data["has_merged_generator_env"]
+        ir.has_private_generator_frame = data["has_private_generator_frame"]
         ir.is_acyclic = data.get("is_acyclic", False)
         ir.is_enum = data["is_enum"]
         ir.coroutine_name = data["is_coroutine"]

mypyc/irbuild/generator.py

@@ -52,6 +52,7 @@
 )
 from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl, gen_generator_func_cleanup
 from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
+from mypyc.irbuild.targets import AssignmentTargetAttr
 from mypyc.primitives.exc_ops import (
     error_catch_op,
     exc_matches_op,
@@ -142,11 +143,7 @@ def instantiate_generator_class(builder: IRBuilder) -> Value:
     fitem = builder.fn_info.fitem
     generator_reg = builder.add(Call(builder.fn_info.generator_class.ir.ctor, [], fitem.line))
 
-    if builder.fn_info.can_merge_generator_and_env_classes():
-        # Set the generator instance to the initial state (zero).
-        zero = Integer(0)
-        builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line))
-    else:
+    if not builder.fn_info.can_merge_generator_and_env_classes():
         # Get the current environment register. If the current function is nested, then the
         # generator class gets instantiated from the callable class' '__call__' method, and hence
         # we use the callable class' environment register. Otherwise, we use the original
@@ -160,9 +157,10 @@ def instantiate_generator_class(builder: IRBuilder) -> Value:
         # defined in the current scope.
         builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line))
 
-        # Set the generator instance's environment to the initial state (zero).
-        zero = Integer(0)
-        builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line))
+    # The continuation label is private generator state even when captured source variables
+    # require a separate environment.
+    zero = Integer(0)
+    builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line))
     return generator_reg
 
 
@@ -171,15 +169,14 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR:
     assert isinstance(builder.fn_info.fitem, FuncDef), builder.fn_info.fitem
     generator_class_ir = mapper.fdef_to_generator[builder.fn_info.fitem]
     generator_class_ir.has_running_flag = True
+    generator_class_ir.has_private_generator_frame = True
     if builder.fn_info.can_merge_generator_and_env_classes():
         builder.fn_info.env_class = generator_class_ir
-        # The merged environment can be thread-confined; see attrs_are_thread_confined.
-        generator_class_ir.has_merged_generator_env = True
     else:
         generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class)
         if not builder.fn_info.fitem.is_coroutine:
-            # After completion generators still need generator.__mypyc_env__ for subsequent
-            # __next__() calls to observe the terminal next-label and raise StopIteration.
+            # The helper currently loads generator.__mypyc_env__ before terminal dispatch, so
+            # an exhausted generator still needs the link on subsequent __next__() calls.
             # Coroutines can't be resumed after completion, so keeping the environment alive
             # there would just extend local lifetimes unnecessarily.
             generator_class_ir.attrs_to_keep_alive_on_completion.add(ENV_ATTR_NAME)
@@ -260,7 +257,9 @@ def add_helper_to_generator_class(
     )
     fn_info.generator_class.ir.methods[GENERATOR_HELPER_NAME] = helper_fn_ir
     builder.functions.append(helper_fn_ir)
-    fn_info.env_class.env_user_function = helper_fn_ir
+    # Compiler-generated values live on the private generator frame even if source-level
+    # captured variables require a separate environment.
+    fn_info.generator_class.ir.env_user_function = helper_fn_ir
 
     return helper_fn_decl
 
@@ -438,12 +437,13 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None:
     else:
         cls.curr_env_reg = load_outer_env(builder, cls.self_reg, builder.symtables[-1])
 
-    # Define a variable representing the label to go to the next time
-    # the '__next__' function of the generator is called, and add it
-    # as an attribute to the environment class.
-    cls.next_label_target = builder.add_var_to_env_class(
-        Var(NEXT_LABEL_ATTR_NAME), int32_rprimitive, cls, reassign=False, always_defined=True
-    )
+    # The continuation label identifies where execution resumes when the generator is next
+    # advanced. Only the serialized generator helper accesses it, so keep it on the private
+    # generator frame instead of a potentially shared closure environment.
+    cls.ir.attributes[NEXT_LABEL_ATTR_NAME] = int32_rprimitive
+    cls.ir.attrs_with_defaults.add(NEXT_LABEL_ATTR_NAME)
+    next_label_target = AssignmentTargetAttr(cls.self_reg, NEXT_LABEL_ATTR_NAME)
+    cls.next_label_target = builder.add_target(Var(NEXT_LABEL_ATTR_NAME), next_label_target)
 
     # Add arguments from the original generator function to the
     # environment of the generator class.

mypyc/test-data/run-generators.test

@@ -763,6 +763,62 @@ def test_basic() -> None:
     assert yields == ('foo',)
     assert val == 3, val
 
+[case testNestedGeneratorPrivateSpill]
+from typing import Generator
+from testutil import run_generator
+
+def make() -> str:
+    return "left"
+
+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)
+
+    return nested("right")
+
+def test_nested_generator_private_spill() -> None:
+    yields, value = run_generator(outer(), ["sent"])
+    assert yields == ("right",)
+    assert value == "leftsent"
+
+[case testGeneratorOverridePrivateSpillsAcrossModules]
+from typing import Generator
+import other_base
+
+def make_number() -> int:
+    return 42
+
+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.
+        # Both generator frames have the same short name, but different full names.
+        return (make_number(), (yield None))[0]
+
+def run(g: Generator[None, None, object]) -> object:
+    assert next(g) is None
+    try:
+        next(g)
+    except StopIteration as e:
+        return e.value
+    assert False
+
+def test_generator_override_private_spills_across_modules() -> None:
+    assert run(other_base.Same().gen()) == "base"
+    assert run(Same().gen()) == 42
+
+[file other_base.py]
+from typing import Generator
+
+def make_text() -> str:
+    return "base"
+
+class Same:
+    def gen(self) -> Generator[None, None, object]:
+        # The call result is live across the yield.
+        return (make_text(), (yield None))[0]
+
 [case testGeneratorReuse]
 from typing import Iterator, Any
 

mypyc/test/test_spill.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import unittest
+
+from mypyc.common import (
+    ENV_ATTR_NAME,
+    GENERATOR_ATTRIBUTE_PREFIX,
+    NEXT_LABEL_ATTR_NAME,
+    TEMP_ATTR_NAME,
+)
+from mypyc.ir.rtypes import RInstance
+from mypyc.test.testutil import build_ir_for_single_file2
+from mypyc.transform.spill import insert_spills
+
+
+class TestSpill(unittest.TestCase):
+    def test_separate_generator_environment_keeps_private_frame_state(self) -> None:
+        # A nested generator needs a separate environment. Since make() is
+        # evaluated before the yield, its result must be spilled across the
+        # suspension point by the post-IRBuild spill pass.
+        source = """\
+def make() -> str:
+    return "left"
+
+def outer():
+    def nested(value: str):
+        return make() + (yield value)
+    return nested("right")
+"""
+        module, _, _, _ = build_ir_for_single_file2(source.splitlines())
+        frame = next(cl for cl in module.classes if cl.has_running_flag)
+        assert frame.env_user_function is not None
+
+        insert_spills(frame.env_user_function, frame)
+
+        env_type = frame.attributes[ENV_ATTR_NAME]
+        assert isinstance(env_type, RInstance)
+        environment = env_type.class_ir
+
+        # Private generator resume state lives on the generator frame, which
+        # 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 + "2_") for name in frame.attributes)
+
+        # Source-level variables stay in the shared environment.
+        assert GENERATOR_ATTRIBUTE_PREFIX + "value" in environment.attributes

mypyc/transform/spill.py

@@ -21,9 +21,10 @@
     SetAttr,
     Value,
 )
+from mypyc.namegen import exported_name
 
 
-def insert_spills(ir: FuncIR, env: ClassIR) -> None:
+def insert_spills(ir: FuncIR, frame: ClassIR) -> None:
     cfg = get_cfg(ir.blocks, use_yields=True)
     live = analyze_live_regs(ir.blocks, cfg)
     entry_live = live.before[ir.blocks[0], 0]
@@ -32,7 +33,7 @@ def insert_spills(ir: FuncIR, env: ClassIR) -> None:
     # TODO: Actually for now, no Registers at all -- we keep the manual spills
     entry_live = {op for op in entry_live if not isinstance(op, Register)}
 
-    ir.blocks = spill_regs(ir.blocks, env, entry_live, live, ir.arg_regs[0])
+    ir.blocks = spill_regs(ir.blocks, frame, entry_live, live, ir.arg_regs[0])
 
 
 def sort_values(values: Collection[Op], blocks: list[BasicBlock]) -> list[Op]:
@@ -50,30 +51,24 @@ def sort_values(values: Collection[Op], blocks: list[BasicBlock]) -> list[Op]:
 
 def spill_regs(
     blocks: list[BasicBlock],
-    env: ClassIR,
+    frame: ClassIR,
     to_spill: set[Value],
     live: AnalysisResult[Value],
-    self_reg: Register,
+    frame_reg: Register,
 ) -> list[BasicBlock]:
-    env_reg: Value
-    for op in blocks[0].ops:
-        if isinstance(op, GetAttr) and op.attr == "__mypyc_env__":
-            env_reg = op
-            break
-    else:
-        # Environment has been merged into generator object
-        env_reg = self_reg
-
     spill_locs = {}
     # Sort values to make the order deterministic. All the spilled values are
     # known to be Op instances, so the cast is safe.
     for i, val in enumerate(sort_values(cast(set[Op], to_spill), blocks)):
-        name = f"{TEMP_ATTR_NAME}2_{i}"
-        env.attributes[name] = val.type
+        # 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}2_{exported_name(frame.fullname)}_{i}"
+        frame.attributes[name] = val.type
         if val.type.error_overlap:
             # We can safely treat as always initialized, since the type has no pointers.
             # This way we also don't need to manage the defined attribute bitfield.
-            env._always_initialized_attrs.add(name)
+            frame._always_initialized_attrs.add(name)
         spill_locs[val] = name
 
     for block in blocks:
@@ -91,13 +86,13 @@ def spill_regs(
                 # value is not live *when we include yields in the
                 # CFG*. (The original decrefs are computed without that.)
                 #
-                # We also skip a decref is the env register is not
+                # We also skip a decref if the frame register is not
                 # live. That should only happen when an exception is
                 # being raised, so everything should be handled there.
-                if op.src not in live.after[block, i] and env_reg in live.after[block, i]:
+                if op.src not in live.after[block, i] and frame_reg in live.after[block, i]:
                     # Skip the DecRef but null out the spilled location
                     null = LoadErrorValue(op.src.type)
-                    block.ops.extend([null, SetAttr(env_reg, spill_locs[op.src], null, op.line)])
+                    block.ops.extend([null, SetAttr(frame_reg, spill_locs[op.src], null, op.line)])
                 continue
 
             if (
@@ -110,7 +105,7 @@ def spill_regs(
                 stolen = op.stolen()
                 for src in op.sources():
                     if src in spill_locs:
-                        read = GetAttr(env_reg, spill_locs[src], op.line)
+                        read = GetAttr(frame_reg, spill_locs[src], op.line)
                         block.ops.append(read)
                         new_sources.append(read)
                         if src.type.is_refcounted and src not in stolen:
@@ -127,6 +122,6 @@ def spill_regs(
 
             if op in spill_locs:
                 # XXX: could we set uninit?
-                block.ops.append(SetAttr(env_reg, spill_locs[op], op, op.line))
+                block.ops.append(SetAttr(frame_reg, spill_locs[op], op, op.line))
 
     return blocks