[mypyc] Avoid polluting generator environment layouts

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)

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

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(
             """\