diff --git a/mypyc/codegen/emitfunc.py b/mypyc/codegen/emitfunc.py index e79faa41570a..35d658aa843d 100644 --- a/mypyc/codegen/emitfunc.py +++ b/mypyc/codegen/emitfunc.py @@ -13,12 +13,12 @@ c_array_initializer, ) from mypyc.common import ( - GENERATOR_ATTRIBUTE_PREFIX, HAVE_IMMORTAL, IS_FREE_THREADED, NATIVE_PREFIX, REG_PREFIX, RUNNING_FIELD, + source_name_from_generator_attribute, ) from mypyc.ir.class_ir import ClassIR from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD, FuncDecl, FuncIR, all_values @@ -485,6 +485,7 @@ def visit_get_attr(self, op: GetAttr) -> None: rtype = op.class_type cl = rtype.class_ir attr_rtype, decl_cl = cl.attr_details(op.attr) + source_attr_name = source_name_from_generator_attribute(op.attr, decl_cl.fullname) prefer_method = cl.is_trait and attr_rtype.error_overlap if cl.get_method(op.attr, prefer_method=prefer_method): # Properties are essentially methods, so use vtable access for them @@ -528,7 +529,7 @@ def visit_get_attr(self, op: GetAttr) -> None: ): # Generate code for the following branch here to avoid # redundant branches in the generated code. - self.emit_attribute_error(branch, cl.name, op.attr) + self.emit_attribute_error(branch, cl.name, source_attr_name) self.emit_line("goto %s;" % self.label(branch.true)) merged_branch = branch self.emitter.emit_line("}") @@ -536,9 +537,7 @@ def visit_get_attr(self, op: GetAttr) -> None: exc_class = "PyExc_AttributeError" self.emitter.emit_line( 'PyErr_SetString({}, "attribute {} of {} undefined");'.format( - exc_class, - repr(op.attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX)), - repr(cl.name), + exc_class, repr(source_attr_name), repr(cl.name) ) ) @@ -1077,7 +1076,7 @@ def emit_traceback(self, op: Branch) -> None: if op.traceback_entry is not None: self.emitter.emit_traceback(self.source_path, self.module_name, op.traceback_entry) - def emit_attribute_error(self, op: Branch, class_name: str, attr: str) -> None: + def emit_attribute_error(self, op: Branch, class_name: str, source_attr_name: str) -> None: assert op.traceback_entry is not None if self.emitter.context.strict_traceback_checks: assert ( @@ -1090,7 +1089,7 @@ def emit_attribute_error(self, op: Branch, class_name: str, attr: str) -> None: self.source_path.replace("\\", "\\\\"), op.traceback_entry[0], class_name, - attr.removeprefix(GENERATOR_ATTRIBUTE_PREFIX), + source_attr_name, op.traceback_entry[1], globals_static, ) diff --git a/mypyc/common.py b/mypyc/common.py index fea09582d5f5..56b0ce1277e7 100644 --- a/mypyc/common.py +++ b/mypyc/common.py @@ -6,6 +6,7 @@ from typing import Any, Final from mypy.util import unnamed_function +from mypyc.namegen import exported_name PREFIX: Final = "CPyPy_" # Python wrappers NATIVE_PREFIX: Final = "CPyDef_" # Native functions etc. @@ -27,8 +28,27 @@ SELF_NAME: Final = "__mypyc_self__" MYPYC_DEFAULTS_SETUP: Final = "__mypyc_defaults_setup" GENERATOR_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_attribute__" +GENERATOR_FRAME_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_frame_attribute__" CPYFUNCTION_NAME = "__cpyfunction__" + +def generator_frame_attribute_prefix(class_fullname: str, *, is_final_class: bool) -> str: + """Return the source-attribute prefix private to a generator frame class.""" + if is_final_class: + return GENERATOR_FRAME_ATTRIBUTE_PREFIX + return f"{GENERATOR_FRAME_ATTRIBUTE_PREFIX}{exported_name(class_fullname)}_" + + +def source_name_from_generator_attribute(name: str, class_fullname: str) -> str: + """Recover a source name from a generator frame or closure attribute.""" + qualified_prefix = generator_frame_attribute_prefix(class_fullname, is_final_class=False) + if name.startswith(qualified_prefix): + return name.removeprefix(qualified_prefix) + if name.startswith(GENERATOR_FRAME_ATTRIBUTE_PREFIX): + return name.removeprefix(GENERATOR_FRAME_ATTRIBUTE_PREFIX) + return name.removeprefix(GENERATOR_ATTRIBUTE_PREFIX) + + # Omits the prefix added to user attribute fields, so it cannot collide with one. RUNNING_FIELD: Final = "mypyc_running" diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index f9964f936d98..0c3deac16ac8 100644 --- a/mypyc/irbuild/builder.py +++ b/mypyc/irbuild/builder.py @@ -27,6 +27,7 @@ DictionaryComprehension, Expression, FuncDef, + FuncItem, GeneratorExpr, IndexExpr, IntExpr, @@ -75,6 +76,7 @@ KEEP_ALIVE_WHOLE_EXPRESSION, MODULE_PREFIX, SELF_NAME, + generator_frame_attribute_prefix, shared_lib_name, ) from mypyc.crash import catch_errors @@ -763,20 +765,23 @@ def get_assignment_target( reg_type = self.type_to_rtype(symbol.type) else: reg_type = self.node_type(lvalue) - # A deleted error-overlap value needs the environment's - # definedness bitmap. Other generator locals start in - # registers and are promoted later if they cross a yield. + # A deleted error-overlap value needs a definedness bitmap. Other generator + # locals start in registers and are promoted later if they cross a yield. if ( self.fn_info.is_generator and reg_type.error_overlap and symbol in self.deleted_vars ): - return self.add_var_to_env_class( - symbol, - reg_type, - self.fn_info.generator_class, - reassign=False, - prefix=GENERATOR_ATTRIBUTE_PREFIX, + if self.is_captured_by_nested_func(symbol): + return self.add_var_to_env_class( + symbol, + reg_type, + self.fn_info.generator_class, + reassign=False, + prefix=GENERATOR_ATTRIBUTE_PREFIX, + ) + return self.add_var_to_generator_frame( + symbol, reg_type, self.fn_info.generator_class.self_reg, reassign=False ) return self.add_local_reg(symbol, reg_type) @@ -1602,6 +1607,44 @@ def add_var_to_env_class( prefix=prefix, ) + def is_free_variable_in_nested_func(self, fitem: FuncItem, symbol: SymbolNode) -> bool: + for nested in self.encapsulating_funcs.get(fitem, []): + if symbol in self.free_variables.get(nested, set()): + return True + if self.is_free_variable_in_nested_func(nested, symbol): + return True + return False + + def is_captured_by_nested_func(self, symbol: SymbolNode) -> bool: + """Does a binding need to be visible to a nested function?""" + return symbol in self.free_variables.get( + self.fn_info.fitem, set() + ) or self.is_free_variable_in_nested_func(self.fn_info.fitem, symbol) + + def add_var_to_generator_frame( + self, + var: SymbolNode, + rtype: RType, + frame_reg: Value, + reassign: bool = False, + always_defined: bool = False, + keep_alive_on_completion: bool = False, + ) -> AssignmentTarget: + """Add a generator-owned source binding to the private generator frame.""" + cls = self.fn_info.generator_class.ir + return self.add_var_to_class( + var, + rtype, + cls, + frame_reg, + reassign=reassign, + always_defined=always_defined, + keep_alive_on_completion=keep_alive_on_completion, + prefix=generator_frame_attribute_prefix( + cls.fullname, is_final_class=cls.is_final_class + ), + ) + def add_var_to_class( self, var: SymbolNode, diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index 529ccd42ef80..4e5c58158cae 100644 --- a/mypyc/irbuild/env_class.py +++ b/mypyc/irbuild/env_class.py @@ -17,7 +17,7 @@ def g() -> int: from __future__ import annotations -from mypy.nodes import Argument, FuncDef, FuncItem, SymbolNode, Var +from mypy.nodes import Argument, FuncDef, SymbolNode, Var from mypyc.common import ( BITMAP_BITS, ENV_ATTR_NAME, @@ -68,21 +68,24 @@ class is generated, the function environment has not yet been return env_class -def finalize_env_class(builder: IRBuilder, prefix: str = "") -> None: +def finalize_env_class(builder: IRBuilder, prefix: str = "", *, add_args: bool = True) -> Value: """Generate, instantiate, and set up the environment of an environment class.""" if not builder.fn_info.can_merge_generator_and_env_classes(): - instantiate_env_class(builder) + env_reg = instantiate_env_class(builder) + else: + env_reg = builder.fn_info.curr_env_reg # Iterate through the function arguments and replace local definitions (using registers) # that were previously added to the environment with references to the function's # environment class. Comprehension scopes have no arguments to add. - if not builder.fn_info.is_comprehension_scope: + if add_args and not builder.fn_info.is_comprehension_scope: if builder.fn_info.is_nested: add_args_to_env( builder, local=False, base=builder.fn_info.callable_class, prefix=prefix ) else: add_args_to_env(builder, local=False, base=builder.fn_info, prefix=prefix) + return env_reg def instantiate_env_class(builder: IRBuilder) -> Value: @@ -105,18 +108,17 @@ def instantiate_env_class(builder: IRBuilder) -> Value: # Top-level functions and comprehension scopes store env reg directly. builder.fn_info._curr_env_reg = curr_env_reg # Comprehension scopes link to parent env if it exists. - if ( - builder.fn_info.is_nested - and builder.fn_infos[-2]._env_class is not None - and builder.fn_infos[-2]._curr_env_reg is not None - ): + parent = builder.fn_infos[-2] if builder.fn_info.is_nested else None + if parent is not None and parent._env_class is not None: + if parent.is_generator: + parent_env_reg = parent.generator_class.curr_env_reg + else: + parent_env_reg = parent._curr_env_reg + else: + parent_env_reg = None + if parent_env_reg is not None: builder.add( - SetAttr( - curr_env_reg, - ENV_ATTR_NAME, - builder.fn_infos[-2].curr_env_reg, - builder.fn_info.fitem.line, - ) + SetAttr(curr_env_reg, ENV_ATTR_NAME, parent_env_reg, builder.fn_info.fitem.line) ) return curr_env_reg @@ -226,8 +228,8 @@ def add_args_to_env( ) -> None: fn_info = builder.fn_info args = fn_info.fitem.arguments - nb = num_bitmap_args(builder, args) if local: + nb = num_bitmap_args(builder, args) for arg in args: rtype = builder.type_to_rtype(arg.variable.type) builder.add_local_reg(arg.variable, rtype, is_arg=True) @@ -235,11 +237,7 @@ def add_args_to_env( builder.add_local_reg(Var(bitmap_name(i)), bitmap_rprimitive, is_arg=True) else: for arg in args: - if ( - is_free_variable(builder, arg.variable) - or fn_info.is_generator - or fn_info.is_coroutine - ): + if is_free_variable(builder, arg.variable): rtype = builder.type_to_rtype(arg.variable.type) assert base is not None, "base cannot be None for adding nonlocal args" builder.add_var_to_env_class( @@ -247,49 +245,60 @@ def add_args_to_env( rtype, base, reassign=reassign, - keep_alive_on_completion=( - is_free_variable(builder, arg.variable) - or is_free_variable_in_nested_func( - builder, builder.fn_info.fitem, arg.variable - ) - ), + keep_alive_on_completion=builder.is_captured_by_nested_func(arg.variable), prefix=prefix, ) +def add_generator_args( + builder: IRBuilder, frame_reg: Value, env_reg: Value, reassign: bool +) -> None: + """Put generator arguments on either the closure environment or private frame.""" + for arg in builder.fn_info.fitem.arguments: + rtype = builder.type_to_rtype(arg.variable.type) + if builder.is_captured_by_nested_func(arg.variable): + builder.add_var_to_class( + arg.variable, + rtype, + builder.fn_info.env_class, + env_reg, + reassign=reassign, + keep_alive_on_completion=True, + prefix=GENERATOR_ATTRIBUTE_PREFIX, + ) + else: + builder.add_var_to_generator_frame(arg.variable, rtype, frame_reg, reassign=reassign) + + def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None: - """Add relevant local variables and nested functions to the environment class. + """Add relevant local variables and nested functions to persistent storage. - Add all variables and functions that are declared/defined within current - function and are referenced in functions nested within this one to this - function's environment class so the nested functions can reference - them even if they are declared after the nested function's definition. - Note that this is done before visiting the body of the function. + Captured bindings go in the environment class. Generator-owned bindings that aren't + captured go on the private generator frame. This is done before visiting the function body + so nested functions can reference declarations that occur later in the body. """ - env_for_func: FuncInfo | ImplicitClass = builder.fn_info - if builder.fn_info.is_generator: - env_for_func = builder.fn_info.generator_class - elif ( - builder.fn_info.is_nested or builder.fn_info.in_non_ext - ) and not builder.fn_info.is_comprehension_scope: - env_for_func = builder.fn_info.callable_class - - if builder.fn_info.fitem in builder.free_variables: + fn_info = builder.fn_info + env_for_func: FuncInfo | ImplicitClass = fn_info + if fn_info.is_generator: + env_for_func = fn_info.generator_class + elif (fn_info.is_nested or fn_info.in_non_ext) and not fn_info.is_comprehension_scope: + env_for_func = fn_info.callable_class + + if fn_info.fitem in builder.free_variables: # Sort the variables to keep things deterministic - for var in sorted(builder.free_variables[builder.fn_info.fitem], key=lambda x: x.name): + for var in sorted(builder.free_variables[fn_info.fitem], key=lambda x: x.name): if isinstance(var, Var): - rtype = builder.type_to_rtype(var.type) builder.add_var_to_env_class( var, - rtype, + builder.type_to_rtype(var.type), env_for_func, reassign=False, keep_alive_on_completion=True, prefix=prefix, ) - if builder.fn_info.fitem in builder.encapsulating_funcs: - for nested_fn in builder.encapsulating_funcs[builder.fn_info.fitem]: + if fn_info.fitem in builder.encapsulating_funcs: + for nested_fn in builder.encapsulating_funcs[fn_info.fitem]: if isinstance(nested_fn, FuncDef): # The return type is 'object' instead of an RInstance of the # callable class because differently defined functions with @@ -299,14 +308,24 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None: nested_prefix = prefix if nested_fn.is_generator or nested_fn.is_coroutine: nested_prefix = GENERATOR_ATTRIBUTE_PREFIX - builder.add_var_to_env_class( - nested_fn, - object_rprimitive, - env_for_func, - reassign=False, - keep_alive_on_completion=is_free_variable(builder, nested_fn), - prefix=nested_prefix, - ) + keep_alive = is_free_variable(builder, nested_fn) + if fn_info.is_generator and not builder.is_captured_by_nested_func(nested_fn): + builder.add_var_to_generator_frame( + nested_fn, + object_rprimitive, + fn_info.generator_class.self_reg, + reassign=False, + keep_alive_on_completion=keep_alive, + ) + else: + builder.add_var_to_env_class( + nested_fn, + object_rprimitive, + env_for_func, + reassign=False, + keep_alive_on_completion=keep_alive, + prefix=nested_prefix, + ) def setup_func_for_recursive_call( @@ -343,14 +362,3 @@ def setup_func_for_recursive_call( def is_free_variable(builder: IRBuilder, symbol: SymbolNode) -> bool: fitem = builder.fn_info.fitem return fitem in builder.free_variables and symbol in builder.free_variables[fitem] - - -def is_free_variable_in_nested_func( - builder: IRBuilder, fitem: FuncItem, symbol: SymbolNode -) -> bool: - for nested in builder.encapsulating_funcs.get(fitem, []): - if symbol in builder.free_variables.get(nested, set()): - return True - if is_free_variable_in_nested_func(builder, nested, symbol): - return True - return False diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index af356928d746..588e1cd51dbb 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py @@ -255,10 +255,9 @@ def c() -> None: generator_class_ir = builder.mapper.fdef_to_generator[fitem] builder.fn_info.generator_class = GeneratorClass(generator_class_ir) - # Functions that contain nested functions need an environment class to store variables that - # are free in their nested functions. Generator functions need an environment class to - # store a variable denoting the next instruction to be executed when the __next__ function - # is called, along with all the variables inside the function itself. + # Functions that contain nested functions need an environment class to store captured + # variables. A generator also needs a separate environment class when it is nested or its + # generated class is non-final. if contains_nested or ( is_generator and not builder.fn_info.can_merge_generator_and_env_classes() ): diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index a3e7dabf6ce2..7cd4f7de2846 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -48,7 +48,7 @@ from mypyc.irbuild.builder import IRBuilder, calculate_arg_defaults, gen_arg_defaults from mypyc.irbuild.context import FuncInfo from mypyc.irbuild.env_class import ( - add_args_to_env, + add_generator_args, add_vars_to_env, finalize_env_class, load_env_registers, @@ -79,10 +79,11 @@ def gen_generator_func( if builder.fn_info.can_merge_generator_and_env_classes(): gen = instantiate_generator_class(builder) builder.fn_info._curr_env_reg = gen - finalize_env_class(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX) + env_reg = finalize_env_class(builder, add_args=False) else: - finalize_env_class(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX) + env_reg = finalize_env_class(builder, add_args=False) gen = instantiate_generator_class(builder) + add_generator_args(builder, gen, env_reg, reassign=True) builder.add(Return(gen)) args, _, blocks, ret_type, fn_info = builder.leave() @@ -179,8 +180,8 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR: else: generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class) if not builder.fn_info.fitem.is_coroutine: - # The helper currently loads generator.__mypyc_env__ before terminal dispatch, so - # an exhausted generator still needs the link on subsequent __next__() calls. + # The helper 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) @@ -436,11 +437,6 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None: cls.stop_iter_value_reg = stop_iter_value_arg cls.self_reg = builder.read(self_target, fitem.line) - if builder.fn_info.can_merge_generator_and_env_classes(): - cls.curr_env_reg = cls.self_reg - else: - 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 @@ -450,11 +446,13 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None: 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. - add_args_to_env( - builder, local=False, base=cls, reassign=False, prefix=GENERATOR_ATTRIBUTE_PREFIX - ) + if builder.fn_info.can_merge_generator_and_env_classes(): + cls.curr_env_reg = cls.self_reg + else: + cls.curr_env_reg = builder.add(GetAttr(cls.self_reg, ENV_ATTR_NAME, fitem.line)) + + # Add arguments from the original generator function to their selected storage objects. + add_generator_args(builder, cls.self_reg, cls.curr_env_reg, reassign=False) # Set the next label register for the generator class. cls.next_label_reg = builder.read(cls.next_label_target, fitem.line) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index d96ec47a3ffe..dc66a375112e 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -49,7 +49,12 @@ YieldExpr, YieldFromExpr, ) -from mypyc.common import GENERATOR_HELPER_NAME, KEEP_ALIVE_SHORT_LIVED, KEEP_ALIVE_WHOLE_EXPRESSION +from mypyc.common import ( + GENERATOR_HELPER_NAME, + KEEP_ALIVE_SHORT_LIVED, + KEEP_ALIVE_WHOLE_EXPRESSION, + source_name_from_generator_attribute, +) from mypyc.ir.ops import ( ERR_NEVER, NAMESPACE_MODULE, @@ -1269,7 +1274,9 @@ def transform_del_item(builder: IRBuilder, target: AssignmentTarget, line: int) if isinstance(target.obj_type, RInstance): cl = target.obj_type.class_ir if not cl.is_deletable(target.attr): - builder.error(f'"{target.attr}" cannot be deleted', line) + _, decl_cl = cl.attr_details(target.attr) + name = source_name_from_generator_attribute(target.attr, decl_cl.fullname) + builder.error(f'"{name}" cannot be deleted', line) builder.note( 'Using "__deletable__ = ' + '[\'\']" in the class body enables "del obj."', diff --git a/mypyc/test-data/irbuild-i64.test b/mypyc/test-data/irbuild-i64.test index 6441e806cdf3..26f8a2f4c495 100644 --- a/mypyc/test-data/irbuild-i64.test +++ b/mypyc/test-data/irbuild-i64.test @@ -2212,7 +2212,7 @@ from mypy_extensions import i64 def gen() -> Iterator[i64]: value: i64 = 1 yield value - # TODO: Improve the error message for generator locals. - del value # E: "__mypyc_generator_attribute__value" cannot be deleted \ + # TODO: Don't suggest __deletable__ for generator locals. + del value # E: "value" cannot be deleted \ # N: Using "__deletable__ = ['']" in the class body enables "del obj." yield 0 diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index f61d7add7048..b881079ab2fd 100644 --- a/mypyc/test-data/run-generators.test +++ b/mypyc/test-data/run-generators.test @@ -806,20 +806,101 @@ def make() -> str: return "left" def outer() -> Generator[str, str, str]: - def nested(value: str) -> Generator[str, str, str]: - for item in [value]: + def nested(captured: str, private: str) -> Generator[str, str, str]: + def callback() -> str: + return captured + + private_local = private + for item in [private_local]: # 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) + # temporaries live across the yield. Only 'captured' stays in the + # separate environment; the other source bindings are frame-private. + return callback() + make() + (yield item) return "unreachable" - return nested("right") + return nested("captured", "right") def test_nested_generator_private_spill() -> None: yields, value = run_generator(outer(), ["sent"]) assert yields == ("right",) - assert value == "leftsent" + assert value == "capturedleftsent" + +[case testGeneratorPrivateAndEscapedBindings] +from typing import Any, Callable, Generator, Iterator + +callbacks: list[Callable[[], int]] = [] + +def gen(shared: list[int], private: int) -> Generator[int, int, None]: + callback = lambda: shared[0] + callbacks.append(callback) + private_local = private + sent = yield private_local + shared = [shared[0] + sent] + yield callback() + +def make_recursive(prefix: str) -> Callable[[int], Generator[object, None, None]]: + def nested(private: int) -> Generator[object, None, None]: + def recurse(count: int) -> str: + if count: + return recurse(count - 1) + return prefix + + yield private + yield recurse(2) + return nested + +def comprehension_lambdas() -> Iterator[list[int]]: + offset = 10 + functions = [lambda: i + offset for i in range(3)] + yield [function() for function in functions] + +def test_private_and_escaped_bindings() -> None: + shared = [1] + g: Any = gen(shared, 7) + assert next(g) == 7 + assert callbacks[-1]() == 1 + shared[0] = 5 + assert callbacks[-1]() == 5 + assert g.send(4) == 9 + assert list(g) == [] + assert callbacks[-1]() == 9 + assert list(make_recursive("outer")(7)) == [7, "outer"] + assert list(comprehension_lambdas()) == [[12, 12, 12]] + +[case testGeneratorOverridePrivateLocals] +from typing import Generator + +def make_text() -> str: + return "base" + +def make_number() -> int: + return 42 + +class Base: + def gen(self) -> Generator[None, None, object]: + # The source local is private and must not alias the differently typed + # local in the derived generator frame. + value = make_text() + yield None + return value + +class Derived(Base): + def gen(self) -> Generator[None, None, object]: + value = make_number() + yield None + return value + +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_locals() -> None: + assert run(Base().gen()) == "base" + assert run(Derived().gen()) == 42 [case testGeneratorOverridePrivateSpillsAcrossModules] from typing import Generator diff --git a/mypyc/test/test_borrow_generator_attrs.py b/mypyc/test/test_borrow_generator_attrs.py index d3c0878c1ec1..7e42eb7ebe94 100644 --- a/mypyc/test/test_borrow_generator_attrs.py +++ b/mypyc/test/test_borrow_generator_attrs.py @@ -8,7 +8,11 @@ import unittest -from mypyc.common import GENERATOR_ATTRIBUTE_PREFIX, TEMP_ATTR_NAME +from mypyc.common import ( + GENERATOR_ATTRIBUTE_PREFIX, + TEMP_ATTR_NAME, + generator_frame_attribute_prefix, +) from mypyc.ir.class_ir import ClassIR from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg from mypyc.ir.ops import ( @@ -257,7 +261,10 @@ def nested() -> str: private_reads = [ op for op in reads - if op.obj is ir.arg_regs[0] and op.attr == GENERATOR_ATTRIBUTE_PREFIX + "other" + if op.obj is ir.arg_regs[0] + and op.attr + == generator_frame_attribute_prefix(cl.fullname, is_final_class=cl.is_final_class) + + "other" ] shared_reads = [ op diff --git a/mypyc/test/test_generator_spills.py b/mypyc/test/test_generator_spills.py index 9e3a0fb76e16..a6e061818692 100644 --- a/mypyc/test/test_generator_spills.py +++ b/mypyc/test/test_generator_spills.py @@ -10,9 +10,11 @@ from mypyc.common import ( ENV_ATTR_NAME, GENERATOR_ATTRIBUTE_PREFIX, + GENERATOR_FRAME_ATTRIBUTE_PREFIX, NEXT_LABEL_ATTR_NAME, SELF_NAME, TEMP_ATTR_NAME, + generator_frame_attribute_prefix, ) from mypyc.ir.class_ir import ClassIR from mypyc.ir.ops import Assign, GetAttr @@ -38,11 +40,8 @@ def promoted_slots(cl: ClassIR) -> set[str]: def frame_variables(cl: ClassIR) -> set[str]: - return { - name.removeprefix(GENERATOR_ATTRIBUTE_PREFIX) - for name in cl.attributes - if name.startswith(GENERATOR_ATTRIBUTE_PREFIX) - } + prefix = generator_frame_attribute_prefix(cl.fullname, is_final_class=cl.is_final_class) + return {name.removeprefix(prefix) for name in cl.attributes if name.startswith(prefix)} class TestGeneratorSpills(unittest.TestCase): @@ -163,6 +162,8 @@ def gen() -> Generator[int, None, int]: """, "gen_gen", ) + assert cl.is_final_class + assert GENERATOR_FRAME_ATTRIBUTE_PREFIX + "crossing" in cl.attributes variables = frame_variables(cl) assert "crossing" in variables assert "local" not in variables diff --git a/mypyc/test/test_spill.py b/mypyc/test/test_spill.py index 9cfa15696bfd..68911d4d3433 100644 --- a/mypyc/test/test_spill.py +++ b/mypyc/test/test_spill.py @@ -7,6 +7,7 @@ GENERATOR_ATTRIBUTE_PREFIX, NEXT_LABEL_ATTR_NAME, TEMP_ATTR_NAME, + generator_frame_attribute_prefix, ) from mypyc.ir.rtypes import RInstance from mypyc.test.testutil import build_ir_for_single_file2 @@ -14,7 +15,7 @@ class TestSpill(unittest.TestCase): - def test_separate_generator_environment_keeps_private_frame_state(self) -> None: + def test_separate_generator_environment_keeps_spills_on_frame(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. @@ -49,5 +50,42 @@ def nested(value: str): assert any(name.startswith(TEMP_ATTR_NAME + "2_") for name in frame.attributes) assert any(name.startswith(TEMP_ATTR_NAME + "3_") for name in frame.attributes) - # Source-level variables stay in the shared environment. - assert GENERATOR_ATTRIBUTE_PREFIX + "value" in environment.attributes + # Noncaptured source-level variables also stay on the private frame. + frame_prefix = generator_frame_attribute_prefix( + frame.fullname, is_final_class=frame.is_final_class + ) + assert frame_prefix + "value" in frame.attributes + assert GENERATOR_ATTRIBUTE_PREFIX + "value" not in environment.attributes + + def test_generator_locals_use_private_frame_unless_captured(self) -> None: + source = """\ +def outer(): + def nested(captured: str, private: str): + def callback() -> str: + return captured + + private_local = private + yield callback() + yield private_local + return nested("captured", "private") +""" + module, _, _, _ = build_ir_for_single_file2(source.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 + frame_attrs = set(frame.attributes) + environment_attrs = set(environment.attributes) + + assert frame.attrs_are_thread_confined() + assert not environment.attrs_are_thread_confined() + + frame_prefix = generator_frame_attribute_prefix( + frame.fullname, is_final_class=frame.is_final_class + ) + assert GENERATOR_ATTRIBUTE_PREFIX + "captured" in environment_attrs + assert frame_prefix + "captured" not in frame_attrs + for name in ("private", "private_local", "callback"): + assert frame_prefix + name in frame_attrs + assert GENERATOR_ATTRIBUTE_PREFIX + name not in environment_attrs diff --git a/mypyc/transform/generator_spills.py b/mypyc/transform/generator_spills.py index 72c0ba22d196..2e42d2dd589f 100644 --- a/mypyc/transform/generator_spills.py +++ b/mypyc/transform/generator_spills.py @@ -3,7 +3,7 @@ from __future__ import annotations from mypyc.analysis.dataflow import analyze_live_regs_with_exception_edges, cleanup_cfg -from mypyc.common import GENERATOR_ATTRIBUTE_PREFIX, TEMP_ATTR_NAME +from mypyc.common import TEMP_ATTR_NAME, generator_frame_attribute_prefix from mypyc.ir.class_ir import ClassIR from mypyc.ir.func_ir import FuncIR from mypyc.ir.ops import ( @@ -114,9 +114,12 @@ def order_registers(registers: list[Register], blocks: list[BasicBlock]) -> list def allocate_slots(frame: ClassIR, registers: list[Register]) -> dict[Register, str]: slots: dict[Register, str] = {} owner = exported_name(frame.fullname) + source_prefix = generator_frame_attribute_prefix( + frame.fullname, is_final_class=frame.is_final_class + ) for index, register in enumerate(registers): if register.name: - name = available_attr_name(frame, GENERATOR_ATTRIBUTE_PREFIX + register.name) + name = available_attr_name(frame, source_prefix + register.name) else: name = f"{TEMP_ATTR_NAME}3_{owner}_{index}" if register.type.error_overlap: