diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 361290e55a1d..548b568621a7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,7 +26,7 @@ repos: hooks: - id: codespell args: - - --ignore-words-list=HAX,Nam,ccompiler,ot,statics,whet,zar + - --ignore-words-list=HAX,Nam,ccompiler,keep-alives,ot,statics,whet,zar exclude: ^(mypy/test/|mypy/typeshed/|mypyc/test-data/|test-data/).+$ - repo: https://github.com/rhysd/actionlint rev: v1.7.7 diff --git a/mypyc/common.py b/mypyc/common.py index 382d640a8408..fa34647c5c72 100644 --- a/mypyc/common.py +++ b/mypyc/common.py @@ -68,6 +68,14 @@ BITMAP_TYPE: Final = "uint32_t" BITMAP_BITS: Final = 32 +# Constant for keeping a (often borrowed) op alive for a short time (a few subexpressions, +# no arbitrary computation or memory allocations), until flush_keep_alives() is called. +KEEP_ALIVE_SHORT_LIVED: Final = 0 +# Keep value alive longer, up to until the current top-level expression has been fully +# evaluated. This allows keeping alive values across function calls and arbitrary +# computation. Note that some expressions (e.g. lambdas), restrict the scope of borrowing. +KEEP_ALIVE_WHOLE_EXPRESSION: Final = 1 + # Runtime C library files that are always included (some ops may bring # extra dependencies via mypyc.ir.deps.SourceDep or mypyc.ir.deps.HeaderDep) RUNTIME_C_FILES: Final = [ diff --git a/mypyc/ir/ops.py b/mypyc/ir/ops.py index 4bc7671b8208..14e12559be1e 100644 --- a/mypyc/ir/ops.py +++ b/mypyc/ir/ops.py @@ -31,7 +31,7 @@ class to enable the new behavior. Sometimes adding a new abstract from mypy_extensions import trait -from mypyc.common import PROPSET_PREFIX +from mypyc.common import KEEP_ALIVE_SHORT_LIVED, PROPSET_PREFIX from mypyc.ir.deps import Dependency from mypyc.ir.rtypes import ( RArray, @@ -898,6 +898,7 @@ def __init__( *, borrow: bool = False, allow_error_value: bool = False, + borrow_scope: int = KEEP_ALIVE_SHORT_LIVED, ) -> None: super().__init__(line) self.obj = obj @@ -912,6 +913,8 @@ def __init__( elif attr_type.error_overlap: self.error_kind = ERR_MAGIC_OVERLAPPING self.is_borrowed = borrow and attr_type.is_refcounted + # How long a borrowed result of this op stays valid (a KEEP_ALIVE_* constant). + self.borrow_scope = borrow_scope def sources(self) -> list[Value]: return [self.obj] diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index 56a3f944fe77..d8ae71e33bf5 100644 --- a/mypyc/irbuild/builder.py +++ b/mypyc/irbuild/builder.py @@ -20,12 +20,17 @@ TYPE_VAR_KIND, TYPE_VAR_TUPLE_KIND, ArgKind, + AssignmentExpr, + AwaitExpr, CallExpr, Decorator, + DictionaryComprehension, Expression, FuncDef, + GeneratorExpr, IndexExpr, IntExpr, + LambdaExpr, Lvalue, MemberExpr, MypyFile, @@ -41,7 +46,10 @@ TypeInfo, TypeParam, Var, + YieldExpr, + YieldFromExpr, ) +from mypy.traverser import TraverserVisitor from mypy.types import ( AnyType, DeletedType, @@ -63,6 +71,8 @@ EXT_SUFFIX, GENERATOR_ATTRIBUTE_PREFIX, IS_FREE_THREADED, + KEEP_ALIVE_SHORT_LIVED, + KEEP_ALIVE_WHOLE_EXPRESSION, MODULE_PREFIX, SELF_NAME, TEMP_ATTR_NAME, @@ -80,6 +90,7 @@ BasicBlock, Branch, Call, + Cast, ComparisonOp, GetAttr, InitStatic, @@ -284,6 +295,20 @@ def __init__( self.imports: dict[str, None] = {} self.can_borrow = False + self.expression_depth = 0 + # Symbols (local vars) reassigned via a walrus expression within the current + # top-level expression. Used to avoid borrowing an attribute over the whole + # expression when the borrow root could be rebound (and thus freed) partway. + self.reassigned_in_expr: set[SymbolNode] = set() + # 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. + self.expr_has_suspend = False + # Saved expression state for enclosing functions (see enter()/leave()). + self.expression_depth_stack: list[int] = [] + self.reassigned_in_expr_stack: list[set[SymbolNode]] = [] + self.expr_has_suspend_stack: list[bool] = [] # When set, load_globals_dict uses this module instead of self.module_name. # Used by generate_attr_defaults_init for cross-module inherited defaults. @@ -315,6 +340,10 @@ def accept(self, node: Statement | Expression, *, can_borrow: bool = False) -> V """ with self.catch_errors(node.line): if isinstance(node, Expression): + self.expression_depth += 1 + if self.expression_depth == 1: + self.reassigned_in_expr = find_walrus_targets(node) + self.expr_has_suspend = expr_has_suspend(node) old_can_borrow = self.can_borrow self.can_borrow = can_borrow try: @@ -329,6 +358,11 @@ def accept(self, node: Statement | Expression, *, can_borrow: bool = False) -> V self.can_borrow = old_can_borrow if not can_borrow: self.flush_keep_alives(node.line) + self.expression_depth -= 1 + if self.expression_depth == 0: + self.flush_keep_alives(node.line, scope=KEEP_ALIVE_WHOLE_EXPRESSION) + self.reassigned_in_expr = set() + self.expr_has_suspend = False return res else: try: @@ -337,8 +371,8 @@ def accept(self, node: Statement | Expression, *, can_borrow: bool = False) -> V pass return None - def flush_keep_alives(self, line: int) -> None: - self.builder.flush_keep_alives(line) + def flush_keep_alives(self, line: int, *, scope: int = KEEP_ALIVE_SHORT_LIVED) -> None: + self.builder.flush_keep_alives(line, scope=scope) # Pass through methods for the most common low-level builder ops, for convenience. @@ -1192,7 +1226,9 @@ def is_synthetic_type(self, typ: TypeInfo) -> bool: return typ.is_named_tuple or typ.is_newtype or typ.typeddict_type is not None def get_final_ref(self, expr: MemberExpr) -> tuple[str, Var, bool] | None: - """Check if `expr` is a final attribute. + """Check if `expr` is a final class or module attribute. + + Return False for instance attributes. This needs to be done differently for class and module attributes to correctly determine fully qualified name. Return a tuple that consists of @@ -1341,6 +1377,15 @@ def enter(self, fn_info: FuncInfo | str = "", *, ret_type: RType = none_rprimiti self.fn_info = fn_info self.fn_infos.append(self.fn_info) self.ret_types.append(ret_type) + # A function body is its own top-level expression context, even when the + # function (e.g. a lambda) is being generated in the middle of an outer + # expression. Save the outer expression state and start fresh. + self.expression_depth_stack.append(self.expression_depth) + self.reassigned_in_expr_stack.append(self.reassigned_in_expr) + self.expr_has_suspend_stack.append(self.expr_has_suspend) + self.expression_depth = 0 + self.reassigned_in_expr = set() + self.expr_has_suspend = False if fn_info.is_generator: self.nonlocal_control.append(GeneratorNonlocalControl()) else: @@ -1354,6 +1399,9 @@ def leave(self) -> tuple[list[Register], list[RuntimeArg], list[BasicBlock], RTy ret_type = self.ret_types.pop() fn_info = self.fn_infos.pop() self.nonlocal_control.pop() + self.expression_depth = self.expression_depth_stack.pop() + self.reassigned_in_expr = self.reassigned_in_expr_stack.pop() + self.expr_has_suspend = self.expr_has_suspend_stack.pop() self.builder = self.builders[-1] self.fn_info = self.fn_infos[-1] return builder.args, runtime_args, builder.blocks, ret_type, fn_info @@ -1389,6 +1437,27 @@ def enter_scope(self, fn_info: FuncInfo) -> Iterator[None]: self.builder = self.builders[-1] self.fn_info = self.fn_infos[-1] + @contextmanager + def enter_borrow_scope(self, line: int) -> Iterator[None]: + """Enter new borrow scope from which borrows can't leak to outer expressions. + + This is a borrow region (see LowLevelIRBuilder.borrow_region) that also + resets the per-expression borrowing heuristic state, since the body forms + its own top-level expression context (e.g. a comprehension iteration or a + lambda body). + """ + old_expression_depth = self.expression_depth + old_reassigned_in_expr = self.reassigned_in_expr + old_expr_has_suspend = self.expr_has_suspend + self.expression_depth = 0 + try: + with self.builder.borrow_region(line): + yield + finally: + self.expression_depth = old_expression_depth + self.reassigned_in_expr = old_reassigned_in_expr + self.expr_has_suspend = old_expr_has_suspend + @contextmanager def enter_method( self, @@ -1581,6 +1650,32 @@ def is_final_native_attr_ref(self, expr: MemberExpr) -> bool: return expr.name in ir.final_attributes return False + def root_is_reassigned(self, v: Value) -> bool: + """Is the root local variable a borrow chain 'v' reads from reassigned this expression? + + A whole-expression borrow of an attribute keeps the borrow root alive only + via the register holding it. If that register belongs to a local variable + that is rebound (via a walrus assignment) during the same top-level + expression, the old value may be freed while the borrow is still live. + """ + if not self.reassigned_in_expr: + return False + # Peel borrowed links back to the root value the chain reads from. + while True: + if isinstance(v, GetAttr) and v.is_borrowed: + v = v.obj + elif isinstance(v, Cast) and v.is_borrowed: + v = v.src + else: + break + if not isinstance(v, Register): + return False + for symbol in self.reassigned_in_expr: + target = self.symtables[-1].get(symbol) + if isinstance(target, AssignmentTargetRegister) and target.register is v: + return True + return False + def mark_block_unreachable(self) -> None: """Mark statements in the innermost block being processed as unreachable. @@ -1695,6 +1790,78 @@ def get_call_target_fullname(ref: RefExpr) -> str: return ref.fullname +class WalrusTargetCollector(TraverserVisitor): + """Collect the symbols assigned to by walrus expressions in a subtree.""" + + def __init__(self) -> None: + self.targets: set[SymbolNode] = set() + + def visit_assignment_expr(self, o: AssignmentExpr) -> None: + if o.target.node is not None: + self.targets.add(o.target.node) + super().visit_assignment_expr(o) + + def visit_lambda_expr(self, o: LambdaExpr) -> None: + # A lambda body forms its own expression context, so don't descend into it. + pass + + +def find_walrus_targets(expr: Expression) -> set[SymbolNode]: + """Return the symbols reassigned via a walrus expression within 'expr'. + + Walrus (':=') is the only way to rebind a variable in the middle of evaluating + an expression, so this is the complete set of in-expression reassignments. + """ + collector = WalrusTargetCollector() + expr.accept(collector) + return collector.targets + + +class SuspendDetector(TraverserVisitor): + """Detect await/yield/yield from expressions in a subtree.""" + + def __init__(self) -> None: + self.found = False + + def visit_await_expr(self, o: AwaitExpr) -> None: + self.found = True + + def visit_yield_expr(self, o: YieldExpr) -> None: + self.found = True + + def visit_yield_from_expr(self, o: YieldFromExpr) -> None: + self.found = True + + def visit_generator_expr(self, o: GeneratorExpr) -> None: + # An 'async for' clause suspends via an implicit await on __anext__ that + # isn't represented as an AwaitExpr node in the AST (list/set comprehensions + # delegate to a GeneratorExpr, so they are covered here too). + if any(o.is_async): + self.found = True + super().visit_generator_expr(o) + + def visit_dictionary_comprehension(self, o: DictionaryComprehension) -> None: + if any(o.is_async): + self.found = True + super().visit_dictionary_comprehension(o) + + def visit_lambda_expr(self, o: LambdaExpr) -> None: + # A lambda body forms its own function (and suspension) context. + pass + + +def expr_has_suspend(expr: Expression) -> bool: + """Does evaluating 'expr' involve a suspension point (await/yield/yield from)? + + A whole-expression borrow can't safely span a suspension point, since the + borrowed value and its borrow root are held in registers that aren't spilled + into the generator environment across the suspend. + """ + detector = SuspendDetector() + expr.accept(detector) + return detector.found + + def create_type_params( builder: IRBuilder, typing_mod: Value, type_args: list[TypeParam], line: int ) -> list[Value]: diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py index b99e8161c08c..91e1e55d8da2 100644 --- a/mypyc/irbuild/expression.py +++ b/mypyc/irbuild/expression.py @@ -58,7 +58,12 @@ TypeType, get_proper_type, ) -from mypyc.common import IS_FREE_THREADED, MAX_SHORT_INT +from mypyc.common import ( + IS_FREE_THREADED, + KEEP_ALIVE_SHORT_LIVED, + KEEP_ALIVE_WHOLE_EXPRESSION, + MAX_SHORT_INT, +) from mypyc.ir.class_ir import ClassIR from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD from mypyc.ir.ops import ( @@ -66,11 +71,14 @@ Assign, BasicBlock, CallC, + Cast, ComparisonOp, + GetAttr, Integer, LoadAddress, LoadLiteral, PrimitiveDescription, + PrimitiveOp, RaiseStandardError, Register, TupleGet, @@ -252,7 +260,7 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value: if expr.fullname in ("typing.TYPE_CHECKING", "typing_extensions.TYPE_CHECKING"): return builder.false(expr.line) - # First check if this is maybe a final attribute. + # First check if this is maybe a final class/module attribute. final = builder.get_final_ref(expr) if final is not None: fullname, final_var, native = final @@ -305,8 +313,43 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value: check_instance_attribute_access_through_class(builder, expr, typ) - borrow = can_borrow and builder.can_borrow - return builder.builder.get_attr(obj, expr.name, rtype, expr.line, borrow=borrow) + is_final = builder.is_final_native_attr_ref(expr) + scope = KEEP_ALIVE_SHORT_LIVED + if ( + is_final + and builder.expression_depth > 1 + and value_borrow_scope(builder, obj) >= KEEP_ALIVE_WHOLE_EXPRESSION + # Don't borrow across the whole expression if the borrow root can be + # rebound via a walrus assignment + and not builder.root_is_reassigned(obj) + # Don't borrow across a suspension point (await/yield/yield from), since + # the borrow would not survive the suspend (registers aren't spilled). + and not builder.expr_has_suspend + ): + scope = KEEP_ALIVE_WHOLE_EXPRESSION + borrow = (can_borrow and builder.can_borrow) or scope == KEEP_ALIVE_WHOLE_EXPRESSION + return builder.builder.get_attr( + obj, expr.name, rtype, expr.line, borrow=borrow, borrow_scope=scope + ) + + +def value_borrow_scope(builder: IRBuilder, v: Value) -> int: + """Compute how long an existing borrowed value can safely be kept alive. + + Returns a KEEP_ALIVE_* constant (or a large value if 'v' is not borrowed and + thus has no borrowing constraint). + """ + if isinstance(v, GetAttr) and v.is_borrowed: + return min(v.borrow_scope, value_borrow_scope(builder, v.obj)) + elif isinstance(v, Cast) and v.is_borrowed: + # A cast just propagates the value (when successful), so the scope is + # determined by the source. + return value_borrow_scope(builder, v.src) + elif isinstance(v, (CallC, PrimitiveOp)) and v.is_borrowed: + # Values borrowed from a C function (e.g. a borrowed list/vec item) may be + # invalidated by arbitrary computation, so they are only short-lived. + return KEEP_ALIVE_SHORT_LIVED + return 999 def check_instance_attribute_access_through_class( @@ -875,15 +918,17 @@ def transform_conditional_expr(builder: IRBuilder, expr: ConditionalExpr) -> Val target = Register(expr_type) builder.activate_block(if_body) - true_value = builder.accept(expr.if_expr) - true_value = builder.coerce(true_value, expr_type, expr.line) - builder.add(Assign(target, true_value, expr.line)) + with builder.builder.borrow_region(expr.line): + true_value = builder.accept(expr.if_expr) + true_value = builder.coerce(true_value, expr_type, expr.line) + builder.add(Assign(target, true_value, expr.line)) builder.goto(next_block) builder.activate_block(else_body) - false_value = builder.accept(expr.else_expr) - false_value = builder.coerce(false_value, expr_type, expr.line) - builder.add(Assign(target, false_value, expr.line)) + with builder.builder.borrow_region(expr.line): + false_value = builder.accept(expr.else_expr) + false_value = builder.coerce(false_value, expr_type, expr.line) + builder.add(Assign(target, false_value, expr.line)) builder.goto(next_block) builder.activate_block(next_block) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index fc539a1fc090..5a36dcb80ff7 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -293,7 +293,8 @@ def sequence_from_generator_preallocate_helper( target_op = empty_op_llbuilder(length, line) def set_item(item_index: Value) -> None: - e = builder.accept(gen.left_expr) + with builder.enter_borrow_scope(line): + e = builder.accept(gen.left_expr) set_item_op(target_op, item_index, e, line) for_loop_helper_with_index( @@ -446,9 +447,10 @@ def loop_contents( remaining_loop_params: the parameters for any further nested loops; if it's empty we'll instead evaluate the "gen_inner_stmts" function """ - # Check conditions, in order, short circuiting them. + # Check conditions, in order, short-circuiting them. for cond in conds: - cond_val = builder.accept(cond) + with builder.enter_borrow_scope(line): + cond_val = builder.accept(cond) cont_block, rest_block = BasicBlock(), BasicBlock() # If the condition is true we'll skip the continue. builder.add_bool_branch(cond_val, rest_block, cont_block) @@ -462,7 +464,8 @@ def loop_contents( else: # We finally reached the actual body of the generator. # Generate the IR for the inner loop body. - gen_inner_stmts() + with builder.enter_borrow_scope(line): + gen_inner_stmts() handle_loop(loop_params) diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py index c0ad1cf1f826..aafdb73fed84 100644 --- a/mypyc/irbuild/ll_builder.py +++ b/mypyc/irbuild/ll_builder.py @@ -7,7 +7,8 @@ from __future__ import annotations import sys -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager from typing import Final, TypeGuard, cast from mypy.argmap import map_actuals_to_formals @@ -19,6 +20,7 @@ FAST_ISINSTANCE_MAX_SUBCLASSES, FAST_PREFIX, IS_FREE_THREADED, + KEEP_ALIVE_SHORT_LIVED, MAX_LITERAL_SHORT_INT, MAX_SHORT_INT, MIN_LITERAL_SHORT_INT, @@ -250,6 +252,15 @@ BOOL_BINARY_OPS: Final = {"&", "&=", "|", "|=", "^", "^=", "==", "!=", "<", "<=", ">", ">="} +class PendingKeepAlive: + __slots__ = ("value", "scope", "seq") + + def __init__(self, value: Value, scope: int, seq: int) -> None: + self.value = value + self.scope = scope + self.seq = seq + + class LowLevelIRBuilder: """A "low-level" IR builder class. @@ -276,7 +287,10 @@ def __init__(self, errors: Errors | None, options: CompilerOptions) -> None: self.error_handlers: list[BasicBlock | None] = [None] # Values that we need to keep alive as long as we have borrowed # temporaries. Use flush_keep_alives() to mark the end of the live range. - self.keep_alives: list[Value] = [] + # The second value is the scope/duration of keep alive (KEEP_ALIVE_* constant). + # Different values must be kept alive for different durations. + self.keep_alives: list[PendingKeepAlive] = [] + self.next_keep_alive_seq = 0 def set_module(self, module_name: str, module_path: str) -> None: """Set the name and path of the current module.""" @@ -367,10 +381,55 @@ def self(self) -> Register: """ return self.args[0] - def flush_keep_alives(self, line: int) -> None: - if self.keep_alives: - self.add(KeepAlive(self.keep_alives.copy(), line)) - self.keep_alives = [] + def flush_keep_alives(self, line: int, *, scope: int = KEEP_ALIVE_SHORT_LIVED) -> None: + if any(entry.scope == scope for entry in self.keep_alives): + self.add( + KeepAlive( + [entry.value for entry in self.keep_alives if entry.scope == scope], line + ) + ) + self.keep_alives = [entry for entry in self.keep_alives if entry.scope != scope] + + def keep_alive_checkpoint(self) -> int: + return self.next_keep_alive_seq + + def flush_keep_alives_since(self, line: int, checkpoint: int) -> None: + """Flush keep-alives added after 'checkpoint' without touching earlier ones.""" + new_keep_alives = [entry for entry in self.keep_alives if entry.seq >= checkpoint] + if new_keep_alives: + self.add(KeepAlive([entry.value for entry in new_keep_alives], line)) + self.keep_alives = [entry for entry in self.keep_alives if entry.seq < checkpoint] + + @contextmanager + def borrow_region(self, line: int) -> Iterator[None]: + """Confine borrows created in this region: flush them when leaving it. + + Keep-alives added inside the region (regardless of their duration scope) + are flushed on exit, so borrows can't leak past a lexical boundary such as + a conditional branch, an 'and'/'or' branch or a comprehension iteration. + """ + checkpoint = self.keep_alive_checkpoint() + yield + self.flush_keep_alives_since(line, checkpoint) + + def add_keep_alive(self, value: Value, scope: int) -> None: + self.keep_alives.append(PendingKeepAlive(value, scope, self.next_keep_alive_seq)) + self.next_keep_alive_seq += 1 + + def keep_borrow_source_alive(self, value: Value, borrow_scope: int) -> None: + """Ensure the source of borrowed value will be alive for given borrow scope. + + A borrow keeps reading from 'value', so 'value' must stay valid for the + borrow's whole scope. When 'value' is itself a borrowed cast it holds no + reference of its own (and coerce() only kept its source alive short-term), + so we follow the cast to the underlying owned value and keep *that* alive + for the requested scope too. Otherwise, a longer-lived borrow through a cast + (e.g. a Final attribute read via cast(T, make())) could read freed memory. + """ + self.add_keep_alive(value, borrow_scope) + while isinstance(value, Cast) and value.is_borrowed: + value = value.src + self.add_keep_alive(value, borrow_scope) def debug_print(self, toprint: str | Value) -> None: if isinstance(toprint, str): @@ -400,7 +459,7 @@ def unbox_or_cast( return self.add(Unbox(src, target_type, line)) else: if can_borrow: - self.keep_alives.append(src) + self.add_keep_alive(src, KEEP_ALIVE_SHORT_LIVED) return self.add(Cast(src, target_type, line, borrow=can_borrow, unchecked=unchecked)) def coerce( @@ -805,19 +864,32 @@ def coerce_nullable(self, src: Value, target_type: RType, line: int) -> Value: # Attribute access def get_attr( - self, obj: Value, attr: str, result_type: RType, line: int, *, borrow: bool = False + self, + obj: Value, + attr: str, + result_type: RType, + line: int, + *, + borrow: bool = False, + borrow_scope: int = KEEP_ALIVE_SHORT_LIVED, ) -> Value: - """Get a native or Python attribute of an object.""" + """Get a native or Python attribute of an object. + + If the result is borrowed, borrow_scope controls how long it stays valid + (a KEEP_ALIVE_* constant). The caller is responsible for choosing a scope + that is actually safe (e.g. downgrading to short-lived if the borrow root + may be rebound during the borrow's live range). + """ if ( isinstance(obj.type, RInstance) and obj.type.class_ir.is_ext_class and obj.type.class_ir.has_attr(attr) ): - op = GetAttr(obj, attr, line, borrow=borrow) + op = GetAttr(obj, attr, line, borrow=borrow, borrow_scope=borrow_scope) # For non-refcounted attribute types, the borrow might be # disabled even if requested, so don't check 'borrow'. if op.is_borrowed: - self.keep_alives.append(obj) + self.keep_borrow_source_alive(obj, borrow_scope) return self.add(op) elif isinstance(obj.type, RUnion): return self.union_get_attr(obj, obj.type, attr, result_type, line) @@ -2112,18 +2184,20 @@ def shortcircuit_helper( # it is the right side if the left is false. true_body, false_body = (right_body, left_body) if op == "and" else (left_body, right_body) - left_value = left() - self.add_bool_branch(left_value, true_body, false_body) + with self.borrow_region(line): + left_value = left() + self.add_bool_branch(left_value, true_body, false_body) - self.activate_block(left_body) - left_coerced = self.coerce(left_value, expr_type, line) - self.add(Assign(target, left_coerced, line)) + self.activate_block(left_body) + left_coerced = self.coerce(left_value, expr_type, line) + self.add(Assign(target, left_coerced, line)) self.goto(next_block) self.activate_block(right_body) - right_value = right() - right_coerced = self.coerce(right_value, expr_type, line) - self.add(Assign(target, right_coerced, line)) + with self.borrow_region(line): + right_value = right() + right_coerced = self.coerce(right_value, expr_type, line) + self.add(Assign(target, right_coerced, line)) self.goto(next_block) self.activate_block(next_block) @@ -2281,7 +2355,7 @@ def call_c( # immediately freed, at the risk of a dangling pointer. for arg in coerced: if not isinstance(arg, (Integer, LoadLiteral)): - self.keep_alives.append(arg) + self.add_keep_alive(arg, KEEP_ALIVE_SHORT_LIVED) if desc.error_kind == ERR_NEG_INT: comp = ComparisonOp(target, Integer(0, desc.return_type, line), ComparisonOp.SGE, line) comp.error_kind = ERR_FALSE @@ -2402,7 +2476,7 @@ def primitive_op( # immediately freed, at the risk of a dangling pointer. for arg in coerced: if not isinstance(arg, (Integer, LoadLiteral)): - self.keep_alives.append(arg) + self.add_keep_alive(arg, KEEP_ALIVE_SHORT_LIVED) if desc.error_kind == ERR_NEG_INT: comp = ComparisonOp(target, Integer(0, desc.return_type, line), ComparisonOp.SGE, line) comp.error_kind = ERR_FALSE diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index e90aa089305f..21f47b190c30 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -49,7 +49,7 @@ YieldExpr, YieldFromExpr, ) -from mypyc.common import TEMP_ATTR_NAME +from mypyc.common import KEEP_ALIVE_SHORT_LIVED, KEEP_ALIVE_WHOLE_EXPRESSION, TEMP_ATTR_NAME from mypyc.ir.ops import ( ERR_NEVER, NAMESPACE_MODULE, @@ -86,7 +86,13 @@ object_rprimitive, ) from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional -from mypyc.irbuild.builder import IRBuilder, create_type_params, int_borrow_friendly_op +from mypyc.irbuild.builder import ( + IRBuilder, + create_type_params, + expr_has_suspend, + find_walrus_targets, + int_borrow_friendly_op, +) from mypyc.irbuild.for_helpers import for_loop_helper from mypyc.irbuild.generator import add_raise_exception_blocks_to_generator_class from mypyc.irbuild.nonlocalcontrol import ( @@ -161,10 +167,17 @@ def transform_expression_stmt(builder: IRBuilder, stmt: ExpressionStmt) -> None: if isinstance(stmt.expr, StrExpr): # Docstring. Ignore return - # ExpressionStmts do not need to be coerced like other Expressions, so we shouldn't - # call builder.accept here. + # ExpressionStmts do not need to be coerced like other Expressions, so + # we shouldn't call builder.accept here. + builder.expression_depth += 1 + builder.reassigned_in_expr = find_walrus_targets(stmt.expr) + builder.expr_has_suspend = expr_has_suspend(stmt.expr) stmt.expr.accept(builder.visitor) - builder.flush_keep_alives(stmt.line) + builder.expression_depth -= 1 + builder.reassigned_in_expr = set() + builder.expr_has_suspend = False + builder.flush_keep_alives(stmt.line, scope=KEEP_ALIVE_SHORT_LIVED) + builder.flush_keep_alives(stmt.line, scope=KEEP_ALIVE_WHOLE_EXPRESSION) def transform_return_stmt(builder: IRBuilder, stmt: ReturnStmt) -> None: diff --git a/mypyc/irbuild/vec.py b/mypyc/irbuild/vec.py index 38615ebe1602..8dffc12a2213 100644 --- a/mypyc/irbuild/vec.py +++ b/mypyc/irbuild/vec.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Final -from mypyc.common import IS_32_BIT_PLATFORM, PLATFORM_SIZE +from mypyc.common import IS_32_BIT_PLATFORM, KEEP_ALIVE_SHORT_LIVED, PLATFORM_SIZE from mypyc.ir.ops import ( ERR_MAGIC, Assign, @@ -351,7 +351,7 @@ def vec_get_item_unsafe_lower( vtype = base.type item_addr = vec_item_ptr(builder, base, index) result = vec_load_mem_item(builder, item_addr, vtype.item_type, can_borrow=can_borrow) - builder.keep_alives.append(base) + builder.add_keep_alive(base, KEEP_ALIVE_SHORT_LIVED) return result diff --git a/mypyc/test-data/irbuild-final.test b/mypyc/test-data/irbuild-final.test new file mode 100644 index 000000000000..0ede988a4320 --- /dev/null +++ b/mypyc/test-data/irbuild-final.test @@ -0,0 +1,252 @@ +[case testFinalAttrBorrowed] +from typing import Final + +class C: + def __init__(self, x: str) -> None: + self.f: Final = x + self.x = x + +def f(s: str) -> None: pass + +def calls(c: C) -> None: + f(c.f) + f(c.x) + c.f.startswith("a") + +def ret(c: C) -> str: + return c.f +[out] +def C.__init__(self, x): + self :: __main__.C + x :: str +L0: + self.f = x + self.x = x + return 1 +def f(s): + s :: str +L0: + return 1 +def calls(c): + c :: __main__.C + r0 :: str + r1 :: None + r2 :: str + r3 :: None + r4, r5 :: str + r6 :: i32 + r7 :: bool +L0: + r0 = borrow c.f + r1 = f(r0) + keep_alive c + r2 = c.x + r3 = f(r2) + r4 = borrow c.f + r5 = 'a' + r6 = CPyStr_Startswith(r4, r5) + r7 = truncate r6: i32 to builtins.bool + keep_alive c + return 1 +def ret(c): + c :: __main__.C + r0 :: str +L0: + r0 = c.f + return r0 + +[case testInheritedFinalAttrBorrowed] +from typing import Final + +class Base: + x: str + + def __init__(self, s: str) -> None: + self.s: Final = s + +class Deriv(Base): + def __init__(self, n: int) -> None: + self.n: Final = n + +def f(b: Base) -> str: + return b.s + "a" + +def g(d: Deriv) -> str: + return d.s * d.n + +def h(d: Deriv) -> str: + return d.x + "a" +[out] +def Base.__init__(self, s): + self :: __main__.Base + s :: str +L0: + self.s = s + return 1 +def Deriv.__init__(self, n): + self :: __main__.Deriv + n :: int +L0: + self.n = n + return 1 +def f(b): + b :: __main__.Base + r0, r1, r2 :: str +L0: + r0 = borrow b.s + r1 = 'a' + r2 = PyUnicode_Concat(r0, r1) + keep_alive b + return r2 +def g(d): + d :: __main__.Deriv + r0 :: str + r1 :: int + r2 :: str +L0: + r0 = borrow d.s + r1 = borrow d.n + r2 = CPyStr_Multiply(r0, r1) + keep_alive d, d + return r2 +def h(d): + d :: __main__.Deriv + r0, r1, r2 :: str +L0: + r0 = d.x + r1 = 'a' + r2 = PyUnicode_Concat(r0, r1) + return r2 + +[case testFinalAttrExpressionKinds] +from typing import Final + +class C: + def __init__(self, s: str, n: int, l: list[str]) -> None: + self.f: Final = s + self.n: Final = n + self.l: Final = l + self.s = s + +def f() -> C: + return C("", 0, []) + +def assign(c: C) -> None: + a = c.f + c.s = c.f + +def indexing(c: C) -> str: + return c.l[c.n] + +def call_eq(c: C) -> bool: + return f().f == c.f +[out] +def C.__init__(self, s, n, l): + self :: __main__.C + s :: str + n :: int + l :: list +L0: + self.f = s + self.n = n + self.l = l + self.s = s + return 1 +def f(): + r0 :: str + r1 :: list + r2 :: __main__.C +L0: + r0 = '' + r1 = PyList_New(0) + r2 = C(r0, 0, r1) + return r2 +def assign(c): + c :: __main__.C + r0, a, r1 :: str + r2 :: bool +L0: + r0 = c.f + a = r0 + r1 = c.f + c.s = r1; r2 = is_error + return 1 +def indexing(c): + c :: __main__.C + r0 :: list + r1 :: int + r2 :: object + r3 :: str +L0: + r0 = borrow c.l + r1 = borrow c.n + r2 = CPyList_GetItem(r0, r1) + r3 = cast(str, r2) + keep_alive c, c + return r3 +def call_eq(c): + c, r0 :: __main__.C + r1, r2 :: str + r3 :: bool +L0: + r0 = f() + r1 = borrow r0.f + r2 = borrow c.f + r3 = CPyStr_Equal(r1, r2) + keep_alive r0, c + return r3 + +[case testFinalAttrExpressionKinds2_withgil] +from typing import Final + +class C: + def __init__(self, d: D) -> None: + self.d: Final = d + +class D: + def __init__(self, s: str) -> None: + self.f: Final = s + +def f(c: C) -> str: + return c.d.f + "a" + +def g(a: list[D], n: int) -> str: + return a[n].f + "a" +[out] +def C.__init__(self, d): + self :: __main__.C + d :: __main__.D +L0: + self.d = d + return 1 +def D.__init__(self, s): + self :: __main__.D + s :: str +L0: + self.f = s + return 1 +def f(c): + c :: __main__.C + r0 :: __main__.D + r1, r2, r3 :: str +L0: + r0 = borrow c.d + r1 = borrow r0.f + r2 = 'a' + r3 = PyUnicode_Concat(r1, r2) + keep_alive c, r0 + return r3 +def g(a, n): + a :: list + n :: int + r0 :: object + r1 :: __main__.D + r2, r3, r4 :: str +L0: + r0 = CPyList_GetItemBorrow(a, n) + r1 = borrow cast(__main__.D, r0) + r2 = r1.f + keep_alive a, n, r0 + r3 = 'a' + r4 = PyUnicode_Concat(r2, r3) + return r4 diff --git a/mypyc/test-data/refcount.test b/mypyc/test-data/refcount.test index 7c7134dcfbfa..bedfabac791d 100644 --- a/mypyc/test-data/refcount.test +++ b/mypyc/test-data/refcount.test @@ -2366,3 +2366,717 @@ L0: r3 = CPyBytes_Concat(r2, c) dec_ref r2 return r3 + +[case testBorrowedFinalAttribute] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(c: C, b: bool) -> str: + a = c.s if b else "x" + return a + +def g(c: C) -> C: + return C(c.s) +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def f(c, b): + c :: __main__.C + b :: bool + r0, r1, r2, a :: str +L0: + if b goto L1 else goto L2 :: bool +L1: + r0 = borrow c.s + inc_ref r0 + r1 = r0 + goto L3 +L2: + r2 = 'x' + inc_ref r2 + r1 = r2 +L3: + a = r1 + return a +def g(c): + c :: __main__.C + r0 :: str + r1 :: __main__.C +L0: + r0 = borrow c.s + r1 = C(r0) + return r1 + +[case testBorrowedFinalAttribute2_withgil] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def g(c: C, n: int) -> str: + a = [c] + return a[n].s + "a" +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def g(c, n): + c :: __main__.C + n :: int + r0 :: list + r1 :: ptr + a :: list + r2 :: object + r3 :: __main__.C + r4, r5, r6 :: str +L0: + r0 = PyList_New(1) + r1 = list_items r0 + inc_ref c + buf_init_item r1, 0, c + a = r0 + r2 = CPyList_GetItemBorrow(a, n) + r3 = borrow cast(__main__.C, r2) + r4 = r3.s + dec_ref a + r5 = 'a' + r6 = PyUnicode_Concat(r4, r5) + dec_ref r4 + return r6 + +[case testBorrowedFinalAttribute3_withgil] +from typing import Final + +class C: + def __init__(self, d: D) -> None: + self.d = d + +class CC: + def __init__(self, d: D) -> None: + self.d: Final = d + +class D: + def __init__(self, s: str) -> None: + self.f: Final = s + +def f() -> str: + c = C(D("x")) + return c.d.f + "a" + +def g(d: D) -> str: + c = CC(d) + return c.d.f + "a" +[out] +def C.__init__(self, d): + self :: __main__.C + d :: __main__.D +L0: + inc_ref d + self.d = d + return 1 +def CC.__init__(self, d): + self :: __main__.CC + d :: __main__.D +L0: + inc_ref d + self.d = d + return 1 +def D.__init__(self, s): + self :: __main__.D + s :: str +L0: + inc_ref s + self.f = s + return 1 +def f(): + r0 :: str + r1 :: __main__.D + r2, c :: __main__.C + r3 :: __main__.D + r4, r5, r6 :: str +L0: + r0 = 'x' + r1 = D(r0) + r2 = C(r1) + dec_ref r1 + c = r2 + r3 = borrow c.d + r4 = r3.f + dec_ref c + r5 = 'a' + r6 = PyUnicode_Concat(r4, r5) + dec_ref r4 + return r6 +def g(d): + d :: __main__.D + r0, c :: __main__.CC + r1 :: __main__.D + r2, r3, r4 :: str +L0: + r0 = CC(d) + c = r0 + r1 = borrow c.d + r2 = borrow r1.f + r3 = 'a' + r4 = PyUnicode_Concat(r2, r3) + dec_ref c + return r4 + +[case testBorrowedFinalAttribute4] +from typing import Final + +class C: + def __init__(self, d: D) -> None: + self.d: Final = d + + def method(self) -> int: + return self.d.foo() + +class D: + def foo(self) -> int: + return 1 +[out] +def C.__init__(self, d): + self :: __main__.C + d :: __main__.D +L0: + inc_ref d + self.d = d + return 1 +def C.method(self): + self :: __main__.C + r0 :: __main__.D + r1 :: int +L0: + r0 = borrow self.d + r1 = r0.foo() + return r1 +def D.foo(self): + self :: __main__.D +L0: + return 2 + +[case testBorrowedFinalAttribute5_64bit] +from typing import Final +from librt.vecs import vec +from mypy_extensions import i64 + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def g(c: C, n: i64) -> str: + a = vec[C]([c]) + return a[n].s + "a" +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def g(c, n): + c :: __main__.C + n :: i64 + r0 :: object + r1 :: ptr + r2 :: vec[__main__.C] + r3, r4 :: ptr + a :: vec[__main__.C] + r5 :: native_int + r6 :: bit + r7 :: i64 + r8 :: bit + r9 :: bool + r10 :: i64 + r11 :: __main__.C + r12, r13, r14 :: str +L0: + r0 = __main__.C :: type + r1 = r0 + r2 = VecTApi.alloc(1, 1, r1) + r3 = r2.items + inc_ref c + set_mem r3, c :: __main__.C* + r4 = r3 + 8 + a = r2 + r5 = a.len + r6 = n < r5 :: unsigned + if r6 goto L4 else goto L1 :: bool +L1: + r7 = n + r5 + r8 = r7 < r5 :: unsigned + if r8 goto L3 else goto L6 :: bool +L2: + r9 = raise IndexError + unreachable +L3: + r10 = r7 + goto L5 +L4: + r10 = n +L5: + r11 = vec_get_item_unsafe_borrow[__main__.C] a, r10 + r12 = r11.s + dec_ref a + r13 = 'a' + r14 = PyUnicode_Concat(r12, r13) + dec_ref r12 + return r14 +L6: + dec_ref a + goto L2 + +[case testBorrowedFinalAttributeReassignViaWalrus] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(c: C, d: C) -> str: + return c.s + (c := d).s +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def f(c, d): + c, d :: __main__.C + r0, r1, r2 :: str +L0: + r0 = c.s + inc_ref d + c = d + dec_ref c + r1 = borrow d.s + r2 = PyUnicode_Concat(r0, r1) + dec_ref r0 + return r2 + +[case testBorrowedFinalAttributeConditionalMerge] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(b: bool) -> str: + return (C("x").s if b else "y") + "z" +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def f(b): + b :: bool + r0 :: str + r1 :: __main__.C + r2, r3, r4, r5, r6 :: str +L0: + if b goto L1 else goto L2 :: bool +L1: + r0 = 'x' + r1 = C(r0) + r2 = borrow r1.s + inc_ref r2 + r3 = r2 + dec_ref r1 + goto L3 +L2: + r4 = 'y' + inc_ref r4 + r3 = r4 +L3: + r5 = 'z' + r6 = PyUnicode_Concat(r3, r5) + dec_ref r3 + return r6 + +[case testBorrowedFinalAttributeAndMerge] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(s: str) -> str: + return (s and C("x").s) + "z" +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def f(s): + s :: str + r0 :: bit + r1, r2 :: str + r3 :: __main__.C + r4, r5, r6 :: str +L0: + r0 = CPyStr_IsTrue(s) + if r0 goto L2 else goto L1 :: bool +L1: + inc_ref s + r1 = s + goto L3 +L2: + r2 = 'x' + r3 = C(r2) + r4 = borrow r3.s + inc_ref r4 + r1 = r4 + dec_ref r3 +L3: + r5 = 'z' + r6 = PyUnicode_Concat(r1, r5) + dec_ref r1 + return r6 + +[case testBorrowedFinalAttributeOrMerge] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(s: str) -> str: + return (s or C("x").s) + "z" +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def f(s): + s :: str + r0 :: bit + r1, r2 :: str + r3 :: __main__.C + r4, r5, r6 :: str +L0: + r0 = CPyStr_IsTrue(s) + if r0 goto L1 else goto L2 :: bool +L1: + inc_ref s + r1 = s + goto L3 +L2: + r2 = 'x' + r3 = C(r2) + r4 = borrow r3.s + inc_ref r4 + r1 = r4 + dec_ref r3 +L3: + r5 = 'z' + r6 = PyUnicode_Concat(r1, r5) + dec_ref r1 + return r6 + +[case testBorrowedFinalAttributeChainedComparisonMerge] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(s: str) -> bool: + return s == "x" == C("x").s +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def f(s): + s, r0 :: str + r1, r2 :: bool + r3 :: str + r4 :: __main__.C + r5 :: str + r6 :: bool +L0: + r0 = 'x' + r1 = CPyStr_EqualLiteral(s, r0, 1) + if r1 goto L2 else goto L1 :: bool +L1: + r2 = r1 + goto L3 +L2: + r3 = 'x' + r4 = C(r3) + r5 = borrow r4.s + r6 = CPyStr_EqualLiteral(r5, r0, 1) + r2 = r6 + dec_ref r4 +L3: + return r2 + +[case testBorrowedFinalInListComprehension_withgil] +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(a: list[C]) -> list[str]: + return [x.s for x in a] +[out] +def C.__init__(self, s): + self :: __main__.C + s :: str +L0: + inc_ref s + self.s = s + return 1 +def f(a): + a :: list + r0 :: native_int + r1 :: list + r2, r3 :: native_int + r4 :: bit + r5 :: object + r6, x :: __main__.C + r7 :: str + r8 :: native_int +L0: + r0 = var_object_size a + r1 = PyList_New(r0) + r2 = 0 +L1: + r3 = var_object_size a + r4 = r2 < r3 :: signed + if r4 goto L2 else goto L4 :: bool +L2: + r5 = list_get_item_unsafe a, r2 + r6 = cast(__main__.C, r5) + x = r6 + r7 = x.s + dec_ref x + CPyList_SetItemUnsafe(r1, r2, r7) +L3: + r8 = r2 + 1 + r2 = r8 + goto L1 +L4: + return r1 + +[case testBorrowedFinalInSetComprehension_withgil] +from typing import Final + +class C: + def __init__(self, d: D) -> None: + self.d: Final = d + +class D: + def __init__(self, s: str) -> None: + self.s: Final = s + +def use(s: str) -> bool: + return True + +def f(a: list[C]) -> set[C]: + return {x for x in a if use(x.d.s)} +[out] +def C.__init__(self, d): + self :: __main__.C + d :: __main__.D +L0: + inc_ref d + self.d = d + return 1 +def D.__init__(self, s): + self :: __main__.D + s :: str +L0: + inc_ref s + self.s = s + return 1 +def use(s): + s :: str +L0: + return 1 +def f(a): + a :: list + r0 :: set + r1, r2 :: native_int + r3 :: bit + r4 :: object + r5, x :: __main__.C + r6 :: __main__.D + r7 :: str + r8 :: bool + r9 :: i32 + r10 :: bit + r11 :: native_int +L0: + r0 = PySet_New(0) + r1 = 0 +L1: + r2 = var_object_size a + r3 = r1 < r2 :: signed + if r3 goto L2 else goto L5 :: bool +L2: + r4 = list_get_item_unsafe a, r1 + r5 = cast(__main__.C, r4) + x = r5 + r6 = borrow x.d + r7 = borrow r6.s + r8 = use(r7) + if r8 goto L3 else goto L6 :: bool +L3: + r9 = PySet_Add(r0, x) + dec_ref x + r10 = r9 >= 0 :: signed +L4: + r11 = r1 + 1 + r1 = r11 + goto L1 +L5: + return r0 +L6: + dec_ref x + goto L4 + +[case testBorrowedFinalAttributeThroughCast] +from typing import Final, cast + +class C: + def __init__(self, x: object) -> None: + self.x: Final = x + +def make() -> object: + return C(1) + +def g(a: object) -> object: + return a + +def h() -> object: + return g(cast(C, make()).x) +[out] +def C.__init__(self, x): + self :: __main__.C + x :: object +L0: + inc_ref x + self.x = x + return 1 +def make(): + r0 :: object + r1 :: __main__.C +L0: + r0 = object 1 + r1 = C(r0) + return r1 +def g(a): + a :: object +L0: + inc_ref a + return a +def h(): + r0 :: object + r1 :: __main__.C + r2, r3 :: object +L0: + r0 = make() + r1 = borrow cast(__main__.C, r0) + r2 = borrow r1.x + r3 = g(r2) + dec_ref r0 + return r3 + +[case testBorrowedFinalAttributeCastChain] +from typing import Final, cast + +class Inner: + def __init__(self, x: object) -> None: + self.x: Final = x + +class Outer: + def __init__(self, inner: object) -> None: + self.inner: Final = inner + +def g(a: object) -> object: + return a + +def make() -> object: + return Outer(Inner(1)) + +def h() -> object: + # cast( Inner, ).x + o = cast(Outer, make()) + return g(cast(Inner, o.inner).x) +[out] +def Inner.__init__(self, x): + self :: __main__.Inner + x :: object +L0: + inc_ref x + self.x = x + return 1 +def Outer.__init__(self, inner): + self :: __main__.Outer + inner :: object +L0: + inc_ref inner + self.inner = inner + return 1 +def g(a): + a :: object +L0: + inc_ref a + return a +def make(): + r0 :: object + r1 :: __main__.Inner + r2 :: __main__.Outer +L0: + r0 = object 1 + r1 = Inner(r0) + r2 = Outer(r1) + dec_ref r1 + return r2 +def h(): + r0 :: object + r1, o :: __main__.Outer + r2 :: object + r3 :: __main__.Inner + r4, r5 :: object +L0: + r0 = make() + r1 = cast(__main__.Outer, r0) + o = r1 + r2 = borrow o.inner + r3 = borrow cast(__main__.Inner, r2) + r4 = borrow r3.x + r5 = g(r4) + dec_ref o + return r5 diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index cf9b54368c27..a7a08e7ced5f 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -1974,3 +1974,96 @@ async def async_iter(vals: list[int]) -> AsyncIterator[int]: yield v [typing fixtures/typing-full.pyi] + +[case testBorrowedFinalAttrAcrossAwait] +from typing import Final +from testutil import async_val + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +async def f(c: C) -> str: + # c.s is a final attr; it must not be borrowed across the await, which is a + # suspension point in the middle of the expression. + return c.s + await async_val("hello") + +[file driver.py] +from native import f, C +from testutil import run_generator + +for i in range(50): + c = C("value" + str(i)) + yields, val = run_generator(f(c), inputs=["world"]) + assert yields == ("hello",), yields + assert val == "value" + str(i) + "world", repr(val) + +[case testBorrowedFinalAttrAcrossAsyncComprehension] +import asyncio +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +class AsyncRange: + def __init__(self, n: int) -> None: + self.n = n + self.i = 0 + + def __aiter__(self) -> "AsyncRange": + return self + + async def __anext__(self) -> str: + if self.i >= self.n: + raise StopAsyncIteration + i = self.i + self.i += 1 + # Suspend in the middle of iteration. + await asyncio.sleep(0) + return "item" + str(i) + +async def f(c: C) -> str: + # c.s is a final attr; it must not be borrowed across the async comprehension, + # which suspends via an implicit await on __anext__ that isn't represented as + # an AwaitExpr node in the AST. + return c.s + "".join([x async for x in AsyncRange(3)]) + +async def test_async_comprehension_borrow() -> None: + for i in range(50): + c = C("value" + str(i)) + assert await f(c) == "value" + str(i) + "item0item1item2" + +[file asyncio/__init__.pyi] +async def sleep(t: float) -> None: ... + +[typing fixtures/typing-full.pyi] + +[case testBorrowedFinalAttrAcrossAwaitAfterComprehension] +import asyncio +from typing import Final + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +async def sink() -> str: + await asyncio.sleep(0) + return "z" + +async def f(c: C, xs: list[int]) -> str: + # A (sync) comprehension enters a nested borrow scope earlier in the same + # top-level expression. It must not clear the enclosing expression's suspend + # flag, or c.s (a final attr) gets borrowed across the await below, which + # produces invalid C (the borrow does not survive the suspension point). + return str([x for x in xs]) + (c.s + await sink()) + +async def test_borrow_final_attr_across_await_after_comprehension() -> None: + for i in range(50): + c = C("v" + str(i)) + assert await f(c, [1, 2, 3]) == "[1, 2, 3]" + "v" + str(i) + "z" + +[file asyncio/__init__.pyi] +async def sleep(t: float) -> None: ... + +[typing fixtures/typing-full.pyi] diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test index 73927bb037a7..4e8b67a0061a 100644 --- a/mypyc/test-data/run-classes.test +++ b/mypyc/test-data/run-classes.test @@ -6104,3 +6104,76 @@ base = sys.getrefcount(BIG) o.v = BIG o.v = BIG assert sys.getrefcount(BIG) == base, "reassignment leaked refs" + +[case testBorrowedFinalAttributeInLambdaAndNestedFunction] +from typing import Final, Callable + +class Box: + def __init__(self, n: int) -> None: + self.n = n + +class C: + def __init__(self, b: Box) -> None: + self.b: Final = b + +def use(x: Box, y: Box) -> int: + return x.n + y.n + +def spend(x: Box) -> Box: + # Allocate a lot so a freed Box slot is likely to be reused, turning any + # use-after-free of a borrowed attribute into an observable wrong result. + junk = [Box(i) for i in range(100)] + return Box(x.n + junk[50].n - 50) + +def make_lambda() -> Callable[[], int]: + # The lambda body creates a C, borrows its final attribute, then calls + # spend() (which allocates heavily) before reading the borrow. + return lambda: use(C(Box(111)).b, spend(Box(0))) + +def make_nested() -> Callable[[], int]: + def inner() -> int: + return use(C(Box(111)).b, spend(Box(0))) + return inner + +def test_borrowed_final_attribute_in_lambda() -> None: + f = make_lambda() + for _ in range(1000): + assert f() == 111 + +def test_borrowed_final_attribute_in_nested_function() -> None: + g = make_nested() + for _ in range(1000): + assert g() == 111 + +[case testBorrowedFinalAttributeInComprehension] +from typing import Final + +class Box: + def __init__(self, n: int) -> None: + self.n = n + +class C: + def __init__(self, b: Box) -> None: + self.b: Final = b + +def use(x: Box, y: Box) -> int: + return x.n + y.n + +def spend(x: Box) -> Box: + # Allocate a lot so a freed Box slot is likely to be reused, turning any + # use-after-free of a borrowed attribute into an observable wrong result. + junk = [Box(i) for i in range(100)] + return Box(x.n + junk[50].n - 50) + +def comp(ns: list[int]) -> list[int]: + # Each iteration allocates a fresh C, borrows its final attribute .b, then + # calls spend() (which allocates heavily) before reading the borrow. The + # borrow container must be kept alive for the current iteration only; if it + # is scoped to the whole (comprehension) expression, the refcount pass + # dec_refs an uninitialized temporary on the loop back-edge and leaks/uses + # freed memory. + return [use(C(Box(111)).b, spend(Box(i))) for i in ns] + +def test_borrowed_final_attribute_in_comprehension() -> None: + for _ in range(1000): + assert comp([1, 2, 3, 4, 5]) == [112, 113, 114, 115, 116] diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index cf1dac7c5733..a23af2ed6eea 100644 --- a/mypyc/test-data/run-generators.test +++ b/mypyc/test-data/run-generators.test @@ -944,3 +944,46 @@ from typing import Optional, Union def test_compiledGeneratorEmptyTuple() -> None: jobs: Generator[Optional[str], None, None] = (_ for _ in ()) assert list(jobs) == [] + +[case testBorrowedFinalAttrAcrossYield] +from typing import Final, Generator + +class C: + def __init__(self, s: str) -> None: + self.s: Final = s + +def f(c: C) -> Generator[str, str, None]: + # c.s is a final attr; it must not be borrowed across the yield, which is + # a suspension point in the middle of the expression. + x = c.s + (yield "first") + yield x + +def test_borrow_across_yield() -> None: + for i in range(50): + c = C("value" + str(i)) + g = f(c) + assert next(g) == "first" + assert g.send("world") == "value" + str(i) + "world" + +def yield_from_gen(c: C) -> Generator[str, str, str]: + r = yield "inner" + return c.s + r + +def g2(c: C) -> Generator[str, str, str]: + # Borrow a final attr across a yield from expression. + x = c.s + (yield from yield_from_gen(c)) + return x + +def test_borrow_across_yield_from() -> None: + for i in range(50): + c = C("v" + str(i)) + gen = g2(c) + assert next(gen) == "inner" + try: + gen.send("!") + except StopIteration as e: + assert e.value == "v" + str(i) + "v" + str(i) + "!", e.value + else: + assert False + +[typing fixtures/typing-full.pyi] diff --git a/mypyc/test/test_irbuild.py b/mypyc/test/test_irbuild.py index 50ffc9004743..6608e6db8e7b 100644 --- a/mypyc/test/test_irbuild.py +++ b/mypyc/test/test_irbuild.py @@ -36,6 +36,7 @@ "irbuild-statements.test", "irbuild-nested.test", "irbuild-classes.test", + "irbuild-final.test", "irbuild-optional.test", "irbuild-any.test", "irbuild-generics.test",