[mypyc] Borrow final attributes more aggressively

.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

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 = [

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]

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]: