diff --git a/mypyc/analysis/dataflow.py b/mypyc/analysis/dataflow.py index 480a725a0aecc..babf7fbc6c87d 100644 --- a/mypyc/analysis/dataflow.py +++ b/mypyc/analysis/dataflow.py @@ -7,6 +7,7 @@ from typing import Any, Generic, TypeVar from mypyc.ir.ops import ( + ERR_NEVER, Assign, AssignMulti, BasicBlock, @@ -518,6 +519,55 @@ def analyze_live_regs(blocks: list[BasicBlock], cfg: CFG) -> AnalysisResult[Valu ) +def analyze_live_regs_with_exception_edges( + blocks: list[BasicBlock], +) -> dict[BasicBlock, set[Value]]: + """Calculate block-entry liveness before exception handling is inserted. + + Unlike get_cfg(), this models an error edge at each operation that can + raise. This matters when an assignment later in the same block overwrites a + value that the error handler can still read. Returns inserted for a yield + are treated as edges to their continuation blocks. + + TODO: Unlike run_analysis(), this supports edges from the middle of basic + blocks, so we reimplement data flow analysis here. Figure out a way + to unify this with run_analysis(). + """ + visitor = LivenessVisitor() + live_in: dict[BasicBlock, set[Value]] = {block: set() for block in blocks} + + while True: + changed = False + for block in reversed(blocks): + terminator = block.terminator + if isinstance(terminator, Return) and terminator.yield_target is not None: + successors = (terminator.yield_target,) + else: + successors = terminator.targets() + + live: set[Value] = set() + for successor in successors: + live.update(live_in[successor]) + + for op in reversed(block.ops): + if ( + block.error_handler is not None + and isinstance(op, RegisterOp) + and op.error_kind != ERR_NEVER + ): + live.update(live_in[block.error_handler]) + gen, kill = op.accept(visitor) + live.difference_update(kill) + live.update(gen) + + if live != live_in[block]: + live_in[block] = live + changed = True + + if not changed: + return live_in + + # Analysis kinds MUST_ANALYSIS = 0 MAYBE_ANALYSIS = 1 diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index ce70dc2eadad6..f9964f936d989 100644 --- a/mypyc/irbuild/builder.py +++ b/mypyc/irbuild/builder.py @@ -75,7 +75,6 @@ KEEP_ALIVE_WHOLE_EXPRESSION, MODULE_PREFIX, SELF_NAME, - TEMP_ATTR_NAME, shared_lib_name, ) from mypyc.crash import catch_errors @@ -152,7 +151,7 @@ AssignmentTargetRegister, AssignmentTargetTuple, ) -from mypyc.irbuild.util import bytes_from_str, is_constant +from mypyc.irbuild.util import bytes_from_str, get_func_def, is_constant from mypyc.irbuild.vec import vec_set_item from mypyc.namegen import exported_name from mypyc.options import CompilerOptions @@ -249,11 +248,8 @@ def __init__( self.callable_class_names: set[str] = set() self.options = options - # These variables keep track of the number of lambdas, implicit indices, and implicit - # iterators instantiated so we avoid name conflicts. The indices and iterators are - # instantiated from for-loops. + # Keep track of the number of lambdas instantiated so we avoid name conflicts. self.lambda_counter = 0 - self.temp_counter = 0 # These variables are populated from the first-pass PreBuildVisitor. self.free_variables = pbv.free_variables @@ -263,6 +259,7 @@ def __init__( self.fdefs_to_decorators = pbv.funcs_to_decorators self.module_import_groups = pbv.module_import_groups self.comprehension_to_fitem = pbv.comprehension_to_fitem + self.deleted_vars = pbv.deleted_vars self.singledispatch_impls = singledispatch_impls @@ -749,11 +746,11 @@ def get_assignment_target( if line == -1: line = lvalue.line if isinstance(lvalue, NameExpr): - # If we are visiting a decorator, then the SymbolNode we really want to be looking at - # is the function that is decorated, not the entire Decorator node itself. + # Use the concrete implementation as the symbol-table key for + # decorated and overloaded functions. symbol = lvalue.node - if isinstance(symbol, Decorator): - symbol = symbol.func + if isinstance(symbol, Decorator | OverloadedFuncDef): + symbol = get_func_def(symbol) if symbol is None: # Semantic analyzer doesn't create ad-hoc Vars for special forms. assert lvalue.is_special_form @@ -766,12 +763,14 @@ def get_assignment_target( reg_type = self.type_to_rtype(symbol.type) else: reg_type = self.node_type(lvalue) - # If the function is a generator function, then first define a new variable - # in the current function's environment class. Next, define a target that - # refers to the newly defined variable in that environment class. Add the - # target to the table containing class environment variables, as well as the - # current environment. - if self.fn_info.is_generator or self.fn_info.is_coroutine: + # 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. + 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, @@ -780,7 +779,6 @@ def get_assignment_target( prefix=GENERATOR_ATTRIBUTE_PREFIX, ) - # Otherwise define a new local variable. return self.add_local_reg(symbol, reg_type) else: # Assign to a previously defined variable. @@ -856,11 +854,6 @@ def read( assert False, "Unsupported lvalue: %r" % target - def read_nullable_attr(self, obj: Value, attr: str, line: int = -1) -> Value: - """Read an attribute that might have an error value without raising AttributeError.""" - assert isinstance(obj.type, RInstance) and obj.type.class_ir.is_ext_class - return self.add(GetAttr(obj, attr, line, allow_error_value=True)) - def assign(self, target: Register | AssignmentTarget, rvalue_reg: Value, line: int) -> None: if isinstance(target, Register): self.add(Assign(target, self.coerce_rvalue(rvalue_reg, target.type, line), line)) @@ -1045,47 +1038,8 @@ def push_loop_stack(self, continue_block: BasicBlock, break_block: BasicBlock) - 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 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_class(Var(name), type, frame.ir, frame.self_reg) - return target - - def spill(self, value: Value) -> AssignmentTarget: - """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) - return target - - def maybe_spill(self, value: Value) -> Value | AssignmentTarget: - """ - 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 - original Value itself for non-generator functions. - """ - if self.fn_info.is_generator: - return self.spill(value) - return value - - def maybe_spill_assignable(self, value: Value) -> Register | AssignmentTarget: - """ - 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 - assignable Register for non-generator functions. - """ - if self.fn_info.is_generator: - return self.spill(value) - + def ensure_register(self, value: Value) -> Register: + """Return an assignable register containing a value.""" if isinstance(value, Register): return value diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py index 6681fcd0825a5..e3b9ad0a2b1a9 100644 --- a/mypyc/irbuild/expression.py +++ b/mypyc/irbuild/expression.py @@ -1415,7 +1415,7 @@ def transform_dictionary_comprehension(builder: IRBuilder, o: DictionaryComprehe def _dict_comp_body(builder: IRBuilder, o: DictionaryComprehension) -> Value: - d = builder.maybe_spill(builder.call_c(dict_new_op, [], o.line)) + d = builder.call_c(dict_new_op, [], o.line) loop_params = list(zip(o.indices, o.sequences, o.condlists, o.is_async)) def gen_inner_stmts() -> None: diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 00315781c6840..f1536fb3111a6 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -321,7 +321,7 @@ def set_item(x: Value, y: Value, z: Value, line: int) -> None: if val is not None: return val - list_ops = builder.maybe_spill(builder.new_list_op([], gen.line)) + list_ops = builder.new_list_op([], gen.line) loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async)) @@ -357,7 +357,7 @@ def translate_set_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value if raise_error_if_contains_unreachable_names(builder, gen): return builder.none() - set_ops = builder.maybe_spill(builder.new_set_op([], gen.line)) + set_ops = builder.new_set_op([], gen.line) loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async)) def gen_inner_stmts() -> None: @@ -687,7 +687,7 @@ def gen_step(self) -> None: def gen_cleanup(self) -> None: """Generate post-loop cleanup (if needed).""" - def load_len(self, expr: Value | AssignmentTarget) -> Value: + def load_len(self, expr: Value) -> Value: """A helper to get collection length, used by several subclasses.""" return self.builder.builder.builtin_len( self.builder.read(expr, self.line), self.line, use_pyssize_t=True @@ -702,13 +702,11 @@ def need_cleanup(self) -> bool: return True 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 - # private generator frame. + # Define a target containing the iterator used by the for-loop. The generator spill + # transform will promote it to the private generator frame if needed. builder = self.builder iter_reg = builder.primitive_op(iter_op, [expr_reg], self.line) - builder.maybe_spill(expr_reg) - self.iter_target = builder.maybe_spill(iter_reg) + self.iter_target = iter_reg self.target_type = target_type def gen_condition(self) -> None: @@ -751,10 +749,9 @@ def need_cleanup(self) -> bool: return True 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 private generator frame. - builder = self.builder - self.iter_target = builder.maybe_spill(expr_reg) + # The generator expression is also the iterator. The generator spill transform will + # promote it to the private generator frame if needed. + self.iter_target = expr_reg self.target_type = target_type def gen_condition(self) -> None: @@ -807,14 +804,11 @@ class ForAsyncIterable(ForGenerator): """Generate IR for an async for loop.""" 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 - # the private generator frame. + # Define a target containing the iterator used by the for-loop. The generator spill + # transform will promote it to the private generator frame if needed. builder = self.builder iter_reg = builder.call_c(aiter_op, [expr_reg], self.line) - builder.maybe_spill(expr_reg) - self.iter_target = builder.maybe_spill(iter_reg) + self.iter_target = iter_reg self.target_type = target_type self.stop_reg = Register(bool_rprimitive) @@ -894,7 +888,7 @@ class ForSequence(ForGenerator): Supports iterating in both forward and reverse. """ - length_reg: Value | AssignmentTarget | None + length_reg: Value | None def init( self, expr_reg: Value, target_type: RType, reverse: bool, length: Value | None = None @@ -907,13 +901,12 @@ def init( # Record a Value indicating the length of the sequence, if known at compile time. self.length = length 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 - # private generator frame. - self.expr_target = builder.maybe_spill(expr_reg) + # The generator spill transform will promote loop state to the private generator frame + # if needed. + self.expr_target = expr_reg if is_immutable_rprimitive(expr_reg.type): # If the expression is an immutable type, we can load the length just once. - self.length_reg = builder.maybe_spill(self.length or self.load_len(self.expr_target)) + self.length_reg = self.length or self.load_len(self.expr_target) else: # Otherwise, even if the length is known, we must recalculate the length # at every iteration for compatibility with python semantics. @@ -926,7 +919,7 @@ def init( else: len_val = self.load_len(self.expr_target) index_reg = builder.builder.int_sub(len_val, 1) - self.index_target = builder.maybe_spill_assignable(index_reg) + self.index_target = builder.ensure_register(index_reg) self.target_type = target_type def gen_condition(self) -> None: @@ -1010,15 +1003,15 @@ def init(self, expr_reg: Value, target_type: RType) -> None: builder = self.builder self.target_type = target_type - # Spill some values so they can be read across yield. - self.expr_target = builder.maybe_spill(expr_reg) + # The generator spill transform will promote loop state that crosses a yield. + self.expr_target = expr_reg offset = Integer(0) - self.offset_target = builder.maybe_spill_assignable(offset) - self.size = builder.maybe_spill(self.load_len(self.expr_target)) + self.offset_target = builder.ensure_register(offset) + self.size = self.load_len(self.expr_target) # For dict class (not a subclass) this is the dictionary itself. iter_reg = builder.call_c(self.dict_iter_op, [expr_reg], self.line) - self.iter_target = builder.maybe_spill(iter_reg) + self.iter_target = iter_reg def gen_condition(self) -> None: """Get next key/value pair, set new offset, and check if we should continue.""" @@ -1132,7 +1125,7 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None: self.start_reg = start_reg self.end_reg = end_reg self.step = step - self.end_target = builder.maybe_spill(end_reg) + self.end_target = end_reg if is_short_int_rprimitive(start_reg.type) and is_short_int_rprimitive(end_reg.type): index_type: RType = short_int_rprimitive elif is_fixed_width_rtype(end_reg.type): @@ -1141,7 +1134,7 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None: index_type = int_rprimitive index_reg = Register(index_type, line=self.line) builder.assign(index_reg, start_reg, self.line) - self.index_reg = builder.maybe_spill_assignable(index_reg) + self.index_reg = index_reg # Initialize loop index to 0. Assert that the index target is assignable. self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index) builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line) @@ -1195,7 +1188,7 @@ def init(self) -> None: # Create a register to store the state of the loop index and # initialize this register along with the loop index to 0. zero = Integer(0) - self.index_reg = builder.maybe_spill_assignable(zero) + self.index_reg = builder.ensure_register(zero) self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index) def gen_step(self) -> None: diff --git a/mypyc/irbuild/main.py b/mypyc/irbuild/main.py index 6614f76e3a825..bc5354c342033 100644 --- a/mypyc/irbuild/main.py +++ b/mypyc/irbuild/main.py @@ -47,6 +47,7 @@ def f(x: int) -> int: from mypyc.irbuild.visitor import IRBuilderVisitor from mypyc.irbuild.vtable import compute_vtable from mypyc.options import CompilerOptions +from mypyc.transform.generator_spills import promote_generator_registers # The stubs for callable contextmanagers are busted so cast it to the # right type... @@ -123,6 +124,13 @@ def build_ir( result[module.fullname] = module_ir class_irs.extend(builder.classes) + # Generator helper calls are fresh C activations after every suspension. + # Move values that must outlive one activation to the private frame before + # attribute-definedness and the remaining IR transforms run. + for class_ir in class_irs: + if class_ir.env_user_function is not None: + promote_generator_registers(class_ir.env_user_function, class_ir) + analyze_always_defined_attrs(class_irs) # Compute vtables. diff --git a/mypyc/irbuild/nonlocalcontrol.py b/mypyc/irbuild/nonlocalcontrol.py index 009a1d6a986bc..60bddcd475035 100644 --- a/mypyc/irbuild/nonlocalcontrol.py +++ b/mypyc/irbuild/nonlocalcontrol.py @@ -12,6 +12,7 @@ from mypyc.ir.ops import ( ERR_NEVER, NO_TRACEBACK_LINE_NO, + Assign, BasicBlock, Branch, Call, @@ -24,7 +25,6 @@ Value, ) from mypyc.ir.rtypes import object_rprimitive -from mypyc.irbuild.targets import AssignmentTarget from mypyc.primitives.exc_ops import restore_exc_info_op, set_stop_iteration_value if TYPE_CHECKING: @@ -93,6 +93,11 @@ class GeneratorNonlocalControl(BaseNonlocalControl): """Default nonlocal control in a generator function outside statements.""" def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None: + # Frame cleanup may clear an attribute holding the return value. Copy it + # first so later spill transforms can place any required load here. + result = Register(value.type) + builder.add(Assign(result, value, line)) + gen_generator_func_cleanup(builder, line) # Raise a StopIteration containing a field for the value that @@ -120,7 +125,7 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None: builder.activate_block(true) # The default/slow path is to raise a StopIteration exception with # return value. - builder.call_c(set_stop_iteration_value, [value], NO_TRACEBACK_LINE_NO) + builder.call_c(set_stop_iteration_value, [result], NO_TRACEBACK_LINE_NO) builder.add(Unreachable()) builder.builder.pop_error_handler() @@ -128,7 +133,7 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None: # The fast path is to store return value via caller-provided pointer # instead of raising an exception. This can only be used when the # caller is a native function. - builder.add(SetMem(object_rprimitive, stop_iter_reg, value)) + builder.add(SetMem(object_rprimitive, stop_iter_reg, result)) builder.add(Return(Integer(0, object_rprimitive))) @@ -178,7 +183,7 @@ class TryFinallyNonlocalControl(NonlocalControl): def __init__(self, target: BasicBlock) -> None: self.target = target - self.ret_reg: None | Register | AssignmentTarget = None + self.ret_reg: Register | None = None def gen_break(self, builder: IRBuilder, line: int) -> None: builder.error("break inside try/finally block is unimplemented", line) @@ -188,13 +193,7 @@ def gen_continue(self, builder: IRBuilder, line: int) -> None: def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None: if self.ret_reg is None: - if builder.fn_info.is_generator: - self.ret_reg = builder.make_spill_target(builder.ret_types[-1]) - else: - self.ret_reg = Register(builder.ret_types[-1]) - # assert needed because of apparent mypy bug... it loses track of the union - # and infers the type as object - assert isinstance(self.ret_reg, (Register, AssignmentTarget)), self.ret_reg + self.ret_reg = Register(builder.ret_types[-1]) builder.assign(self.ret_reg, value, line) builder.add(Goto(self.target)) @@ -207,7 +206,7 @@ class ExceptNonlocalControl(CleanupNonlocalControl): This is super annoying. """ - def __init__(self, outer: NonlocalControl, saved: Value | AssignmentTarget) -> None: + def __init__(self, outer: NonlocalControl, saved: Value) -> None: super().__init__(outer) self.saved = saved diff --git a/mypyc/irbuild/prebuildvisitor.py b/mypyc/irbuild/prebuildvisitor.py index 232fcc55a0a77..2d9535f06a52c 100644 --- a/mypyc/irbuild/prebuildvisitor.py +++ b/mypyc/irbuild/prebuildvisitor.py @@ -4,6 +4,7 @@ AssignmentStmt, Block, Decorator, + DelStmt, DictionaryComprehension, Expression, FuncDef, @@ -11,12 +12,14 @@ GeneratorExpr, Import, LambdaExpr, + ListExpr, MemberExpr, MypyFile, NameExpr, Node, StrExpr, SymbolNode, + TupleExpr, Var, ) from mypy.traverser import ExtendedTraverserVisitor, TraverserVisitor @@ -94,6 +97,10 @@ def __init__( # All property setters encountered so far. self.prop_setters: set[FuncDef] = set() + # Local variables that occur in a del statement. This is needed before + # their first assignment is lowered (see IRBuilder.get_assignment_target). + self.deleted_vars: set[Var] = set() + # A map from any function that contains nested functions to # a set of all the functions that are nested within it. self.encapsulating_funcs: dict[FuncItem, list[FuncItem]] = {} @@ -147,6 +154,17 @@ def visit_block(self, block: Block) -> None: super().visit_block(block) self._current_import_group = None + def visit_del_stmt(self, stmt: DelStmt) -> None: + self.record_deleted_names(stmt.expr) + super().visit_del_stmt(stmt) + + def record_deleted_names(self, expr: Expression) -> None: + if isinstance(expr, NameExpr) and isinstance(expr.node, Var): + self.deleted_vars.add(expr.node) + elif isinstance(expr, (ListExpr, TupleExpr)): + for item in expr.items: + self.record_deleted_names(item) + def visit_decorator(self, dec: Decorator) -> None: if dec.decorators: # Only add the function being decorated if there exist diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 7ebfc2b4ca29e..a3e571b633776 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -49,12 +49,7 @@ YieldExpr, YieldFromExpr, ) -from mypyc.common import ( - GENERATOR_HELPER_NAME, - KEEP_ALIVE_SHORT_LIVED, - KEEP_ALIVE_WHOLE_EXPRESSION, - TEMP_ATTR_NAME, -) +from mypyc.common import GENERATOR_HELPER_NAME, KEEP_ALIVE_SHORT_LIVED, KEEP_ALIVE_WHOLE_EXPRESSION from mypyc.ir.ops import ( ERR_NEVER, NAMESPACE_MODULE, @@ -716,7 +711,7 @@ def transform_try_except( # exception is raised, based on the exception in exc_info. builder.builder.push_error_handler(double_except_block) builder.activate_block(except_entry) - old_exc = builder.maybe_spill(builder.call_c(error_catch_op, [], line)) + old_exc = builder.call_c(error_catch_op, [], line) # Compile the except blocks with the nonlocal control flow overridden to clear exc_info builder.nonlocal_control.append(ExceptNonlocalControl(builder.nonlocal_control[-1], old_exc)) @@ -796,7 +791,7 @@ def try_finally_try( return_entry: BasicBlock, main_entry: BasicBlock, try_body: GenFunc, -) -> Register | AssignmentTarget | None: +) -> Register | None: # Compile the try block with an error handler control = TryFinallyNonlocalControl(return_entry) builder.builder.push_error_handler(err_handler) @@ -817,7 +812,7 @@ def try_finally_entry_blocks( return_entry: BasicBlock, main_entry: BasicBlock, finally_block: BasicBlock, - ret_reg: Register | AssignmentTarget | None, + ret_reg: Register | None, ) -> Value: line = builder.fn_info.fitem.line old_exc = Register(exc_rtuple, line=line) @@ -862,7 +857,7 @@ def try_finally_resolve_control( cleanup_block: BasicBlock, finally_control: FinallyNonlocalControl, old_exc: Value, - ret_reg: Register | AssignmentTarget | None, + ret_reg: Register | None, ) -> BasicBlock: """Resolve the control flow out of a finally block. @@ -883,15 +878,10 @@ def try_finally_resolve_control( if ret_reg: builder.activate_block(rest) return_block, rest = BasicBlock(), BasicBlock() - # For spill targets in try/finally, use nullable read to avoid AttributeError - if isinstance(ret_reg, AssignmentTargetAttr) and ret_reg.attr.startswith(TEMP_ATTR_NAME): - ret_val = builder.read_nullable_attr(ret_reg.obj, ret_reg.attr, line) - else: - ret_val = builder.read(ret_reg, line) - builder.add(Branch(ret_val, rest, return_block, Branch.IS_ERROR, line)) + builder.add(Branch(ret_reg, rest, return_block, Branch.IS_ERROR, line)) builder.activate_block(return_block) - builder.nonlocal_control[-1].gen_return(builder, ret_val, line) + builder.nonlocal_control[-1].gen_return(builder, ret_reg, line) # TODO: handle break/continue builder.activate_block(rest) @@ -1136,11 +1126,11 @@ def transform_with( exit_ = None else: typ = builder.primitive_op(type_op, [mgr_v], line) - exit_ = builder.maybe_spill(builder.py_get_attr(typ, f"__{al}exit__", line)) + exit_ = builder.py_get_attr(typ, f"__{al}exit__", line) value = builder.py_call(builder.py_get_attr(typ, f"__{al}enter__", line), [mgr_v], line) - mgr = builder.maybe_spill(mgr_v) - exc = builder.maybe_spill_assignable(builder.true()) + mgr = mgr_v + exc = builder.ensure_register(builder.true()) if is_async: value = emit_await(builder, value, line) @@ -1213,7 +1203,7 @@ def transform_with_lock( # __enter__: acquire the lock value = builder.primitive_op(lock_acquire_op, [mgr_v], line) - mgr = builder.maybe_spill(mgr_v) + mgr = mgr_v def try_body() -> None: if target: @@ -1288,7 +1278,7 @@ def transform_del_item(builder: IRBuilder, target: AssignmentTarget, line: int) # Delete a local by assigning an error value to it, which will # prompt the insertion of uninit checks. builder.add( - Assign(target.register, builder.add(LoadErrorValue(target.type, undefines=True))) + Assign(target.register, builder.add(LoadErrorValue(target.type, undefines=True)), line) ) elif isinstance(target, AssignmentTargetTuple): for subtarget in target.items: @@ -1311,7 +1301,7 @@ def emit_yield(builder: IRBuilder, val: Value, line: int) -> Value: next_label = len(cls.continuation_blocks) cls.continuation_blocks.append(next_block) builder.assign(cls.next_label_target, Integer(next_label), line) - builder.add(Return(retval, yield_target=next_block)) + builder.add(Return(retval, line, yield_target=next_block)) builder.activate_block(next_block) add_raise_exception_blocks_to_generator_class(builder, line) @@ -1349,7 +1339,7 @@ def emit_yield_from_or_await( else: iter_val = builder.call_c(get_op, [val], line) - iter_reg = builder.maybe_spill_assignable(iter_val) + iter_reg = builder.ensure_register(iter_val) stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock() diff --git a/mypyc/test-data/irbuild-i64.test b/mypyc/test-data/irbuild-i64.test index 4cea7894052aa..6441e806cdf33 100644 --- a/mypyc/test-data/irbuild-i64.test +++ b/mypyc/test-data/irbuild-i64.test @@ -2203,3 +2203,16 @@ L3: r10 = (1, r6) t = r10 return 1 + +[case testDeleteI64LocalInGenerator] +from typing import Iterator + +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 \ + # N: Using "__deletable__ = ['']" in the class body enables "del obj." + yield 0 diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 971902494d9d5..9a729f8fff461 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -135,6 +135,50 @@ async def sleep(t: float) -> None: ... [typing fixtures/typing-full.pyi] +[case testRunAsyncLocalsAcrossAwait] +import asyncio +from typing import List + +from testutil import assertRaises + +async def increment(value: int) -> int: + await asyncio.sleep(0) + return value + 1 + +async def accumulate(count: int) -> int: + total = 0 + for value in range(count): + total += await increment(value) + return total + +async def return_through_finally(log: List[str]) -> int: + try: + return 7 + finally: + await asyncio.sleep(0) + log.append("finally") + +async def deleted_across_await(delete: bool) -> str: + value = "live" + if delete: + del value + await asyncio.sleep(0) + return value + +async def test_async_locals_across_await() -> None: + assert await accumulate(4) == 10 + log: List[str] = [] + assert await return_through_finally(log) == 7 + assert log == ["finally"] + assert await deleted_across_await(False) == "live" + with assertRaises(AttributeError): + await deleted_across_await(True) + +[file asyncio/__init__.pyi] +async def sleep(delay: float) -> None: ... + +[typing fixtures/typing-full.pyi] + [case testRunAsyncAwaitInVariousPositions] from typing import cast, Any diff --git a/mypyc/test-data/run-floats.test b/mypyc/test-data/run-floats.test index 424d52cdb0d55..08a5fbbb12c78 100644 --- a/mypyc/test-data/run-floats.test +++ b/mypyc/test-data/run-floats.test @@ -543,3 +543,16 @@ class subc(float): [file userdefinedfloat.py] class float: pass + +[case testFloatErrorValueAcrossYield] +from typing import Iterator + +from testutil import FLOAT_MAGIC + +def gen() -> Iterator[float]: + value = FLOAT_MAGIC + yield 0.0 + yield value + +def test_float_error_value_across_yield() -> None: + assert list(gen()) == [0.0, FLOAT_MAGIC] diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index 67c569bf0b23c..437b384c1fa43 100644 --- a/mypyc/test-data/run-generators.test +++ b/mypyc/test-data/run-generators.test @@ -1139,6 +1139,97 @@ def test_generator_method_override_still_works() -> None: with assertRaises(ValueError, "generator already executing"): next(g) +[case testExceptionSurvivesYieldInFinally] +from typing import Iterator +from testutil import assertRaises + +class Marker(Exception): + pass + +def gen() -> Iterator[int]: + try: + raise Marker + finally: + yield 1 + +def test_exception_survives_yield_in_finally() -> None: + it = gen() + assert next(it) == 1 + with assertRaises(Marker): + next(it) + +[case testCrossingLocalOnExceptionOnlyPath] +from typing import Iterator + +class InnerError(Exception): + pass + +class AssignmentError(Exception): + pass + +def fail_assignment() -> str: + raise AssignmentError + +def through_handler() -> Iterator[str]: + value = "old handler value" + try: + try: + yield "first" + raise InnerError + except InnerError: + value = fail_assignment() + except AssignmentError: + yield value + +def through_finally() -> Iterator[str]: + value = "old finally value" + try: + try: + yield "first" + raise InnerError + finally: + value = fail_assignment() + except AssignmentError: + yield value + +def test_crossing_local_on_exception_only_path() -> None: + for it, expected in [ + (through_handler(), "old handler value"), + (through_finally(), "old finally value"), + ]: + assert next(it) == "first" + assert next(it) == expected + +[case testDelLocalInGenerator] +from typing import Iterator + +from testutil import assertRaises + +def crossing(delete: bool) -> Iterator[str]: + value = "crossing" + yield value + if delete: + del value + yield "pause" + yield value + +def non_crossing(delete: bool) -> Iterator[str]: + value = "local" + if delete: + del value + yield value + +def test_del_local_in_generator() -> None: + assert list(crossing(False)) == ["crossing", "pause", "crossing"] + assert list(non_crossing(False)) == ["local"] + it = crossing(True) + assert next(it) == "crossing" + assert next(it) == "pause" + with assertRaises(AttributeError, "attribute 'value' of 'crossing_gen' undefined"): + next(it) + with assertRaises(UnboundLocalError, 'local variable "value" referenced before assignment'): + list(non_crossing(True)) + [case testYieldFromNativeGeneratorInInterpretedCode] from typing import Generator diff --git a/mypyc/test-data/run-i64.test b/mypyc/test-data/run-i64.test index 1a2d6fbbecfbe..59630e65cd67f 100644 --- a/mypyc/test-data/run-i64.test +++ b/mypyc/test-data/run-i64.test @@ -1556,3 +1556,28 @@ def test_dunder_arg_check() -> None: 2**63 in o with assertRaises(IndexError): 1 in o + +[case testI64DefinednessAcrossYield] +from typing import Final, Iterator + +from mypy_extensions import i64 +from testutil import assertRaises + +MAGIC: Final = -113 + +def error_value() -> Iterator[i64]: + value: i64 = MAGIC + yield 0 + yield value + +def conditional(define: bool) -> Iterator[i64]: + if define: + value: i64 = 5 + yield 0 + yield value + +def test_i64_definedness_across_yield() -> None: + assert list(error_value()) == [0, MAGIC] + assert list(conditional(True)) == [0, 5] + with assertRaises(AttributeError): + list(conditional(False)) diff --git a/mypyc/test/test_generator_spills.py b/mypyc/test/test_generator_spills.py new file mode 100644 index 0000000000000..440ebb068675e --- /dev/null +++ b/mypyc/test/test_generator_spills.py @@ -0,0 +1,194 @@ +"""Tests for promoting generator registers onto their frames. + +These don't use the normal data-driven tests, since generators produce very verbose IR. +""" + +from __future__ import annotations + +import unittest + +from mypyc.common import GENERATOR_ATTRIBUTE_PREFIX, TEMP_ATTR_NAME +from mypyc.ir.class_ir import ClassIR +from mypyc.ir.ops import Assign, GetAttr +from mypyc.test.testutil import build_ir_for_single_file2 +from mypyc.transform.generator_spills import registers_live_across_yield + + +def generator_class(source: str, name: str) -> ClassIR: + module, _, _, _ = build_ir_for_single_file2(source.splitlines()) + matches = [ + cl + for cl in module.classes + if (cl.name == name or cl.name.startswith(name + "___")) + and cl.env_user_function is not None + ] + assert len(matches) == 1 + return matches[0] + + +def promoted_slots(cl: ClassIR) -> set[str]: + return {name for name in cl.attributes if name.startswith(TEMP_ATTR_NAME + "3_")} + + +def frame_variables(cl: ClassIR) -> set[str]: + return { + name.removeprefix(GENERATOR_ATTRIBUTE_PREFIX) + for name in cl.attributes + if name.startswith(GENERATOR_ATTRIBUTE_PREFIX) + } + + +class TestGeneratorSpills(unittest.TestCase): + def test_saved_exception_state_is_promoted(self) -> None: + cl = generator_class( + """\ +from typing import Iterator + +def gen() -> Iterator[int]: + try: + yield 1 + finally: + yield 2 +""", + "gen_gen", + ) + helper = cl.env_user_function + assert helper is not None + slots = promoted_slots(cl) + assert len(slots) == 1 + reads = [ + op + for block in helper.blocks + for op in block.ops + if isinstance(op, GetAttr) and op.attr in slots + ] + assert reads and all(read.allow_error_value for read in reads) + assert not registers_live_across_yield(helper) + + def test_address_taken_register_is_not_promoted(self) -> None: + cl = generator_class( + """\ +from typing import Any, Generator + +def gen(values: Any) -> Generator[int, None, None]: + yield from values +""", + "gen_gen", + ) + helper = cl.env_user_function + assert helper is not None + assert not registers_live_across_yield(helper) + + def test_loop_state_only_crossing_a_yield_uses_a_slot(self) -> None: + source = """\ +from typing import Iterator + +def before_yield(values: list[int]) -> Iterator[int]: + total = 0 + for value in values: + total += value + yield total + +def inside_loop(values: list[int]) -> Iterator[int]: + for value in values: + yield value +""" + before = generator_class(source, "before_yield_gen") + inside = generator_class(source, "inside_loop_gen") + before_slots = promoted_slots(before) + inside_slots = promoted_slots(inside) + assert not before_slots + assert len(inside_slots) == 1 + + def test_return_value_is_copied_before_frame_cleanup(self) -> None: + cl = generator_class( + """\ +from typing import Any + +async def collect(values: Any) -> Any: + return {value: value async for value in values} +""", + "collect_gen", + ) + helper = cl.env_user_function + assert helper is not None + result_copies = [ + op + for block in helper.blocks + for op in block.ops + if isinstance(op, Assign) and not op.dest.is_arg + ] + assert result_copies + + def test_only_crossing_named_local_is_on_frame(self) -> None: + cl = generator_class( + """\ +from typing import Generator + +def gen() -> Generator[int, None, int]: + local = 1 + crossing = 2 + yield local + return crossing +""", + "gen_gen", + ) + variables = frame_variables(cl) + assert "crossing" in variables + assert "local" not in variables + + def test_captured_local_stays_in_environment(self) -> None: + module, _, _, _ = build_ir_for_single_file2("""\ +from typing import Any, Iterator + +def gen() -> Iterator[Any]: + captured = "value" + def get() -> str: + return captured + yield get +""".splitlines()) + frame = next(cl for cl in module.classes if cl.name == "gen_gen") + environment = next( + cl for cl in module.classes if GENERATOR_ATTRIBUTE_PREFIX + "captured" in cl.attributes + ) + assert environment is not frame + + def test_nested_finally_promotes_each_saved_exception(self) -> None: + cl = generator_class( + """\ +from typing import Iterator + +def gen() -> Iterator[int]: + try: + try: + yield 1 + finally: + yield 2 + finally: + yield 3 +""", + "gen_gen", + ) + helper = cl.env_user_function + assert helper is not None + slots = promoted_slots(cl) + assert len(slots) == 2 + assert not registers_live_across_yield(helper) + + def test_recursive_generator_reference_is_rematerialized(self) -> None: + cl = generator_class( + """\ +from typing import Any, Iterator + +def outer(edges: list[list[int]]) -> Any: + def walk(node: int) -> Iterator[int]: + for child in edges[node]: + yield from walk(child) + yield node + return walk +""", + "walk_gen", + ) + # The recursive reference comes from the closure again on every resume, + # so it doesn't need a second copy on the generator frame. + assert "walk" not in frame_variables(cl) diff --git a/mypyc/test/test_spill.py b/mypyc/test/test_spill.py index 9e2831b196c90..9cfa15696bfd1 100644 --- a/mypyc/test/test_spill.py +++ b/mypyc/test/test_spill.py @@ -24,11 +24,12 @@ def make() -> str: def outer(): def nested(value: str): + result = "" 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" + # The loop iterator is promoted from a register during IR building, + # while the result of make() is spilled by the later transform. + result += make() + (yield item) + return result return nested("right") """ module, _, _, _ = build_ir_for_single_file2(source.splitlines()) @@ -45,8 +46,8 @@ 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) + 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 diff --git a/mypyc/transform/generator_spills.py b/mypyc/transform/generator_spills.py new file mode 100644 index 0000000000000..72c0ba22d1964 --- /dev/null +++ b/mypyc/transform/generator_spills.py @@ -0,0 +1,185 @@ +"""Move generator registers that survive suspension onto the generator frame.""" + +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.ir.class_ir import ClassIR +from mypyc.ir.func_ir import FuncIR +from mypyc.ir.ops import ( + NO_TRACEBACK_LINE_NO, + Assign, + AssignMulti, + BaseAssign, + BasicBlock, + Branch, + GetAttr, + LoadAddress, + LoadErrorValue, + Op, + Register, + SetAttr, + Value, +) +from mypyc.namegen import exported_name + + +def promote_generator_registers(ir: FuncIR, frame: ClassIR) -> None: + """Replace registers live across a suspension with frame attributes. + + This runs before uninitialized-value, exception, and refcount transforms. + """ + cleanup_cfg(ir.blocks) + registers = registers_live_across_yield(ir) + if not registers: + return + + nullable = nullable_registers(ir.blocks) + slots = allocate_slots(frame, registers) + rewrite_registers(ir.blocks, ir.arg_regs[0], slots, nullable) + + # Keep this close to the rewrite: it catches a new op kind that is not + # handled above without restating the liveness analysis. + for block in ir.blocks: + for op in block.ops: + assert not any(source in slots for source in op.sources()) + assert not (isinstance(op, BaseAssign) and op.dest in slots) + + +def registers_live_across_yield(ir: FuncIR) -> list[Register]: + """Return promotable registers live at generator-helper entry.""" + live_in = analyze_live_regs_with_exception_edges(ir.blocks) + excluded = unpromotable_registers(ir.blocks) + registers = [ + value + for value in live_in[ir.blocks[0]] + if isinstance(value, Register) and not value.is_arg and value not in excluded + ] + return order_registers(registers, ir.blocks) + + +def unpromotable_registers(blocks: list[BasicBlock]) -> set[Register]: + """Find registers whose representation requires a real C local.""" + result: set[Register] = set() + for block in blocks: + for op in block.ops: + if isinstance(op, LoadAddress) and isinstance(op.src, Register): + result.add(op.src) + elif isinstance(op, AssignMulti): + result.add(op.dest) + elif ( + isinstance(op, Assign) + and isinstance(op.src, LoadErrorValue) + and op.src.undefines + and op.dest.type.error_overlap + ): + result.add(op.dest) + return result + + +def nullable_registers(blocks: list[BasicBlock]) -> set[Register]: + """Find registers for which the error value is an ordinary state.""" + result: set[Register] = set() + for block in blocks: + for op in block.ops: + if ( + isinstance(op, Branch) + and op.op == Branch.IS_ERROR + and isinstance(op.value, Register) + ): + result.add(op.value) + elif ( + isinstance(op, Assign) + and isinstance(op.src, LoadErrorValue) + and not op.src.undefines + ): + result.add(op.dest) + return result + + +def order_registers(registers: list[Register], blocks: list[BasicBlock]) -> list[Register]: + """Order registers by first appearance for deterministic frame layouts.""" + positions: dict[Register, int] = {} + for block in blocks: + for op in block.ops: + values = op.sources() + if isinstance(op, BaseAssign): + values = [op.dest, *values] + for value in values: + if isinstance(value, Register) and value not in positions: + positions[value] = len(positions) + return sorted(registers, key=positions.__getitem__) + + +def allocate_slots(frame: ClassIR, registers: list[Register]) -> dict[Register, str]: + slots: dict[Register, str] = {} + owner = exported_name(frame.fullname) + for index, register in enumerate(registers): + if register.name: + name = available_attr_name(frame, GENERATOR_ATTRIBUTE_PREFIX + register.name) + else: + name = f"{TEMP_ATTR_NAME}3_{owner}_{index}" + if register.type.error_overlap: + # Compiler temporaries are assigned before use. Recording a + # default avoids allocating an undefinedness bitmap for them. + frame.attrs_with_defaults.add(name) + frame.attributes[name] = register.type + slots[register] = name + return slots + + +def available_attr_name(frame: ClassIR, base: str) -> str: + if not frame.has_attr(base): + return base + suffix = 2 + while frame.has_attr(f"{base}__{suffix}"): + suffix += 1 + return f"{base}__{suffix}" + + +def rewrite_registers( + blocks: list[BasicBlock], + frame_reg: Register, + slots: dict[Register, str], + nullable: set[Register], +) -> None: + """Replace promoted register accesses with generator frame attribute accesses.""" + # IRTransform maps each value to one replacement, but mutable registers need a fresh + # attribute read at each use, so we can't use IRTransform here. + for block in blocks: + old_ops = block.ops + block.ops = [] + for op in old_ops: + replacements: dict[Register, GetAttr] = {} + new_sources: list[Value] = [] + for source in op.sources(): + if isinstance(source, Register) and source in slots: + read = replacements.get(source) + if read is None: + read = GetAttr( + frame_reg, + slots[source], + source_line(op, source), + allow_error_value=source in nullable, + ) + block.ops.append(read) + replacements[source] = read + new_sources.append(read) + else: + new_sources.append(source) + op.set_sources(new_sources) + + if isinstance(op, Assign) and op.dest in slots: + block.ops.append( + SetAttr(frame_reg, slots[op.dest], op.src, source_line(op, op.dest)) + ) + else: + block.ops.append(op) + + +def source_line(op: Op, register: Register) -> int: + if op.line >= 0 or op.line == NO_TRACEBACK_LINE_NO: + return op.line + if register.line >= 0: + return register.line + return NO_TRACEBACK_LINE_NO diff --git a/mypyc/transform/spill.py b/mypyc/transform/spill.py index ee07edf24328f..313e795262131 100644 --- a/mypyc/transform/spill.py +++ b/mypyc/transform/spill.py @@ -1,9 +1,11 @@ -"""Insert spills for values that are live across yields.""" +"""Spill operation results that are live across generator suspension points. + +Registers are spilled in an earlier pass. +""" from __future__ import annotations from collections.abc import Collection -from typing import cast from mypyc.analysis.dataflow import AnalysisResult, analyze_live_regs, get_cfg from mypyc.common import TEMP_ATTR_NAME @@ -27,13 +29,14 @@ 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] + # Registers needing frame storage were handled earlier by promote_generator_registers(). + # Argument, address-taken, and RArray registers stay as C locals. - entry_live = {op for op in entry_live if not (isinstance(op, Register) and op.is_arg)} - # 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)} + # Each generator helper invocation starts at the resume dispatch, so an operation result + # live at entry must have survived a previous suspension. + to_spill = {value for value in live.before[ir.blocks[0], 0] if isinstance(value, Op)} - ir.blocks = spill_regs(ir.blocks, frame, entry_live, live, ir.arg_regs[0]) + ir.blocks = spill_regs(ir.blocks, frame, to_spill, live, ir.arg_regs[0]) def sort_values(values: Collection[Op], blocks: list[BasicBlock]) -> list[Op]: @@ -52,14 +55,13 @@ def sort_values(values: Collection[Op], blocks: list[BasicBlock]) -> list[Op]: def spill_regs( blocks: list[BasicBlock], frame: ClassIR, - to_spill: set[Value], + to_spill: set[Op], live: AnalysisResult[Value], frame_reg: Register, ) -> list[BasicBlock]: 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)): + # Sort values to make the order deterministic. + for i, val in enumerate(sort_values(to_spill, blocks)): # 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.