diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index 1f8bc773d03e..ce70dc2eadad 100644 --- a/mypyc/irbuild/builder.py +++ b/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,14 +1046,18 @@ 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) @@ -1061,7 +1065,7 @@ def spill(self, value: Value) -> AssignmentTarget: 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: diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 5a36dcb80ff7..cecc69a07e27 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/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) diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index b4c33ea60b05..8e09fb881310 100644 --- a/mypyc/test-data/run-generators.test +++ b/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 diff --git a/mypyc/test/test_spill.py b/mypyc/test/test_spill.py index 898f6e927529..9e2831b196c9 100644 --- a/mypyc/test/test_spill.py +++ b/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.