From dcc3b52bafbb38e31b0d3574cdf0f4be8562136d Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Tue, 15 Sep 2026 13:00:45 +0100 Subject: [PATCH] [mypyc] Avoid polluting generator environment layouts When setting up a non-merged generator, we passed the generator helper symbol table to load_outer_env. This made the closure environment class inherit the helper arguments and generator state as attributes, adding six unused object slots to every affected environment. It could also replace the declared type of the environment self-reference with the generator frame type. Load the outer environment directly from the generator frame during setup, and make load_outer_env copy only attributes that are actually owned by the environment being loaded. Add a regression test that verifies helper-only attributes stay out of the environment and that its self-reference keeps the correct type. --- mypyc/irbuild/env_class.py | 25 ++++++++++++---------- mypyc/irbuild/generator.py | 5 +++-- mypyc/test/test_generator_spills.py | 32 ++++++++++++++++++++++++++++- 3 files changed, 48 insertions(+), 14 deletions(-) diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index 286d7bfbf1079..529ccd42ef80c 100644 --- a/mypyc/irbuild/env_class.py +++ b/mypyc/irbuild/env_class.py @@ -149,16 +149,14 @@ def load_outer_env( *, borrow: bool = False, ) -> Value: - """Load the environment class for a given base into a register. + """Load the environment referenced by "base" and bind its symbols. - Additionally, iterates through all of the SymbolNode and - AssignmentTarget instances of the environment at the given index's - symtable, and adds those instances to the environment of the - current environment. This is done so that the current environment - can access outer environment variables without having to reload - all of the environment registers. + "outer_env" may contain registers and attributes from multiple environment levels. + Bind only attributes stored directly on the loaded environment, since the other targets + cannot be accessed through it. - Returns the register where the environment class was loaded. + If "borrow" is true, return a borrowed reference. In that case the environment link on + "base" must be final so that the reference remains valid. """ if borrow: assert isinstance(base.type, RInstance) @@ -167,9 +165,14 @@ def load_outer_env( assert isinstance(env.type, RInstance), f"{env} must be of type RInstance" for symbol, target in outer_env.items(): - attr_name = symbol.name - if isinstance(target, AssignmentTargetAttr): - attr_name = target.attr + # Only targets actually stored in this environment are visible through it. + if ( + not isinstance(target, AssignmentTargetAttr) + or not isinstance(target.obj.type, RInstance) + or target.obj.type.class_ir is not env.type.class_ir + ): + continue + attr_name = target.attr env.type.class_ir.attributes[attr_name] = target.type symbol_target = AssignmentTargetAttr(env, attr_name) builder.add_target(symbol, symbol_target) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 40791e8326f8d..a3e7dabf6ce20 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -26,6 +26,7 @@ BasicBlock, Branch, Call, + GetAttr, Goto, Integer, LoadErrorValue, @@ -51,7 +52,6 @@ add_vars_to_env, finalize_env_class, load_env_registers, - load_outer_env, load_outer_envs, setup_func_for_recursive_call, ) @@ -439,7 +439,8 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None: if builder.fn_info.can_merge_generator_and_env_classes(): cls.curr_env_reg = cls.self_reg else: - cls.curr_env_reg = load_outer_env(builder, cls.self_reg, builder.symtables[-1]) + cls.curr_env_reg = builder.add(GetAttr(cls.self_reg, ENV_ATTR_NAME, fitem.line)) + assert isinstance(cls.curr_env_reg.type, RInstance) # 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 diff --git a/mypyc/test/test_generator_spills.py b/mypyc/test/test_generator_spills.py index 440ebb068675e..9e3a0fb76e164 100644 --- a/mypyc/test/test_generator_spills.py +++ b/mypyc/test/test_generator_spills.py @@ -7,9 +7,16 @@ import unittest -from mypyc.common import GENERATOR_ATTRIBUTE_PREFIX, TEMP_ATTR_NAME +from mypyc.common import ( + ENV_ATTR_NAME, + GENERATOR_ATTRIBUTE_PREFIX, + NEXT_LABEL_ATTR_NAME, + SELF_NAME, + TEMP_ATTR_NAME, +) from mypyc.ir.class_ir import ClassIR from mypyc.ir.ops import Assign, GetAttr +from mypyc.ir.rtypes import RInstance from mypyc.test.testutil import build_ir_for_single_file2 from mypyc.transform.generator_spills import registers_live_across_yield @@ -39,6 +46,29 @@ def frame_variables(cl: ClassIR) -> set[str]: class TestGeneratorSpills(unittest.TestCase): + def test_generator_environment_excludes_helper_arguments(self) -> None: + # Loading the outer environment must not add the generator helper's arguments or + # state to the environment class, or overwrite the type of its self-reference. + module, _, _, _ = build_ir_for_single_file2("""\ +from typing import Iterator + +def outer() -> Iterator[str]: + captured = "value" + def nested() -> str: + return captured + yield nested() +""".splitlines()) + frame = next(cl for cl in module.classes if cl.has_running_flag) + env_type = frame.attributes[ENV_ATTR_NAME] + assert isinstance(env_type, RInstance) + environment = env_type.class_ir + + self_type = environment.attributes[SELF_NAME] + assert isinstance(self_type, RInstance) + assert self_type.class_ir is environment + for name in ("type", "value", "traceback", "arg", "stop_iter_ptr", NEXT_LABEL_ATTR_NAME): + assert name not in environment.attributes + def test_saved_exception_state_is_promoted(self) -> None: cl = generator_class( """\