[mypyc] Borrow operands of several primitives (#12810)

Borrow an operand such as `x.y` (attribute of a native class) in various contexts 
when it's safe to do so. This reduces the number of incref/decref operations we 
need to perform. This continues work started in #12805.

These cases now support borrowing (for some subexpressions, in some contexts):
* `x.y is None`
* Cast source value
* `len(x.y)` (if the operand is a list)
* `isinstance(x.y, C)`
* `x.y[a.b]`
* `x.y.z = 1`

Author: JukkaL

mypyc/ir/ops.py

@@ -775,15 +775,18 @@ class Cast(RegisterOp):
 
     error_kind = ERR_MAGIC
 
-    def __init__(self, src: Value, typ: RType, line: int) -> None:
+    def __init__(self, src: Value, typ: RType, line: int, *, borrow: bool = False) -> None:
         super().__init__(line)
         self.src = src
         self.type = typ
+        self.is_borrowed = borrow
 
     def sources(self) -> List[Value]:
         return [self.src]
 
     def stolen(self) -> List[Value]:
+        if self.is_borrowed:
+            return []
         return [self.src]
 
     def accept(self, visitor: 'OpVisitor[T]') -> T:

mypyc/ir/pprint.py

@@ -77,11 +77,12 @@ def visit_load_literal(self, op: LoadLiteral) -> str:
         return self.format('%r = %s%s', op, prefix, repr(op.value))
 
     def visit_get_attr(self, op: GetAttr) -> str:
+        return self.format('%r = %s%r.%s', op, self.borrow_prefix(op), op.obj, op.attr)
+
+    def borrow_prefix(self, op: Op) -> str:
         if op.is_borrowed:
-            borrow = 'borrow '
-        else:
-            borrow = ''
-        return self.format('%r = %s%r.%s', op, borrow, op.obj, op.attr)
+            return 'borrow '
+        return ''
 
     def visit_set_attr(self, op: SetAttr) -> str:
         if op.is_init:
@@ -142,7 +143,7 @@ def visit_method_call(self, op: MethodCall) -> str:
         return s
 
     def visit_cast(self, op: Cast) -> str:
-        return self.format('%r = cast(%s, %r)', op, op.type, op.src)
+        return self.format('%r = %scast(%s, %r)', op, self.borrow_prefix(op), op.type, op.src)
 
     def visit_box(self, op: Box) -> str:
         return self.format('%r = box(%s, %r)', op, op.src.type, op.src)

mypyc/irbuild/builder.py

@@ -182,7 +182,7 @@ def accept(self, node: Union[Statement, Expression], *,
                     res = Register(self.node_type(node))
                 self.can_borrow = old_can_borrow
                 if not can_borrow:
-                    self.builder.flush_keep_alives()
+                    self.flush_keep_alives()
                 return res
             else:
                 try:
@@ -191,6 +191,9 @@ def accept(self, node: Union[Statement, Expression], *,
                     pass
                 return None
 
+    def flush_keep_alives(self) -> None:
+        self.builder.flush_keep_alives()
+
     # Pass through methods for the most common low-level builder ops, for convenience.
 
     def add(self, op: Op) -> Value:
@@ -234,7 +237,7 @@ def binary_op(self, lreg: Value, rreg: Value, expr_op: str, line: int) -> Value:
         return self.builder.binary_op(lreg, rreg, expr_op, line)
 
     def coerce(self, src: Value, target_type: RType, line: int, force: bool = False) -> Value:
-        return self.builder.coerce(src, target_type, line, force)
+        return self.builder.coerce(src, target_type, line, force, can_borrow=self.can_borrow)
 
     def none_object(self) -> Value:
         return self.builder.none_object()
@@ -510,7 +513,8 @@ def get_assignment_target(self, lvalue: Lvalue,
             return AssignmentTargetIndex(base, index)
         elif isinstance(lvalue, MemberExpr):
             # Attribute assignment x.y = e
-            obj = self.accept(lvalue.expr)
+            can_borrow = self.is_native_attr_ref(lvalue)
+            obj = self.accept(lvalue.expr, can_borrow=can_borrow)
             return AssignmentTargetAttr(obj, lvalue.name)
         elif isinstance(lvalue, TupleExpr):
             # Multiple assignment a, ..., b = e
@@ -1176,6 +1180,14 @@ def load_module_attr_by_fullname(self, fullname: str, line: int) -> Value:
         left = self.load_module(module)
         return self.py_get_attr(left, name, line)
 
+    def is_native_attr_ref(self, expr: MemberExpr) -> bool:
+        """Is expr a direct reference to a native (struct) attribute of an instance?"""
+        obj_rtype = self.node_type(expr.expr)
+        return (isinstance(obj_rtype, RInstance)
+                and obj_rtype.class_ir.is_ext_class
+                and obj_rtype.class_ir.has_attr(expr.name)
+                and not obj_rtype.class_ir.get_method(expr.name))
+
     # Lacks a good type because there wasn't a reasonable type in 3.5 :(
     def catch_errors(self, line: int) -> Any:
         return catch_errors(self.module_path, line)

mypyc/irbuild/expression.py

@@ -21,7 +21,8 @@
     Value, Register, TupleGet, TupleSet, BasicBlock, Assign, LoadAddress, RaiseStandardError
 )
 from mypyc.ir.rtypes import (
-    RTuple, RInstance, object_rprimitive, is_none_rprimitive, int_rprimitive, is_int_rprimitive
+    RTuple, object_rprimitive, is_none_rprimitive, int_rprimitive, is_int_rprimitive,
+    is_list_rprimitive
 )
 from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD
 from mypyc.irbuild.format_str_tokenizer import (
@@ -130,17 +131,8 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
     if isinstance(expr.node, MypyFile) and expr.node.fullname in builder.imports:
         return builder.load_module(expr.node.fullname)
 
-    obj_rtype = builder.node_type(expr.expr)
-    if (isinstance(obj_rtype, RInstance)
-            and obj_rtype.class_ir.is_ext_class
-            and obj_rtype.class_ir.has_attr(expr.name)
-            and not obj_rtype.class_ir.get_method(expr.name)):
-        # Direct attribute access -> can borrow object
-        can_borrow = True
-    else:
-        can_borrow = False
+    can_borrow = builder.is_native_attr_ref(expr)
     obj = builder.accept(expr.expr, can_borrow=can_borrow)
-
     rtype = builder.node_type(expr)
 
     # Special case: for named tuples transform attribute access to faster index access.
@@ -418,8 +410,12 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
 
 
 def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
-    base = builder.accept(expr.base)
     index = expr.index
+    base_type = builder.node_type(expr.base)
+    is_list = is_list_rprimitive(base_type)
+    can_borrow_base = is_list and is_borrow_friendly_expr(builder, index)
+
+    base = builder.accept(expr.base, can_borrow=can_borrow_base)
 
     if isinstance(base.type, RTuple) and isinstance(index, IntExpr):
         return builder.add(TupleGet(base, index.value, expr.line))
@@ -429,11 +425,31 @@ def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
         if value:
             return value
 
-    index_reg = builder.accept(expr.index)
+    index_reg = builder.accept(expr.index, can_borrow=is_list)
     return builder.gen_method_call(
         base, '__getitem__', [index_reg], builder.node_type(expr), expr.line)
 
 
+def is_borrow_friendly_expr(builder: IRBuilder, expr: Expression) -> bool:
+    """Can the result of the expression borrowed temporarily?
+
+    Borrowing means keeping a reference without incrementing the reference count.
+    """
+    if isinstance(expr, (IntExpr, FloatExpr, StrExpr, BytesExpr)):
+        # Literals are immportal and can always be borrowed
+        return True
+    if isinstance(expr, (UnaryExpr, OpExpr)) and constant_fold_expr(builder, expr) is not None:
+        # Literal expressions are similar to literals
+        return True
+    if isinstance(expr, NameExpr):
+        if isinstance(expr.node, Var) and expr.kind == LDEF:
+            # Local variable reference can be borrowed
+            return True
+    if isinstance(expr, MemberExpr) and builder.is_native_attr_ref(expr):
+        return True
+    return False
+
+
 def try_constant_fold(builder: IRBuilder, expr: Expression) -> Optional[Value]:
     """Return the constant value of an expression if possible.
 
@@ -513,7 +529,8 @@ def transform_conditional_expr(builder: IRBuilder, expr: ConditionalExpr) -> Val
 def transform_comparison_expr(builder: IRBuilder, e: ComparisonExpr) -> Value:
     # x in (...)/[...]
     # x not in (...)/[...]
-    if (e.operators[0] in ['in', 'not in']
+    first_op = e.operators[0]
+    if (first_op in ['in', 'not in']
             and len(e.operators) == 1
             and isinstance(e.operands[1], (TupleExpr, ListExpr))):
         items = e.operands[1].items
@@ -560,6 +577,12 @@ def transform_comparison_expr(builder: IRBuilder, e: ComparisonExpr) -> Value:
             else:
                 return builder.true()
 
+    if first_op in ('is', 'is not') and len(e.operators) == 1:
+        right = e.operands[1]
+        if isinstance(right, NameExpr) and right.fullname == 'builtins.None':
+            # Special case 'is None' / 'is not None'.
+            return translate_is_none(builder, e.operands[0], negated=first_op != 'is')
+
     # TODO: Don't produce an expression when used in conditional context
     # All of the trickiness here is due to support for chained conditionals
     # (`e1 < e2 > e3`, etc). `e1 < e2 > e3` is approximately equivalent to
@@ -584,6 +607,11 @@ def go(i: int, prev: Value) -> Value:
     return go(0, builder.accept(e.operands[0]))
 
 
+def translate_is_none(builder: IRBuilder, expr: Expression, negated: bool) -> Value:
+    v = builder.accept(expr, can_borrow=True)
+    return builder.binary_op(v, builder.none_object(), 'is not' if negated else 'is', expr.line)
+
+
 def transform_basic_comparison(builder: IRBuilder,
                                op: str,
                                left: Value,

mypyc/irbuild/ll_builder.py

@@ -163,13 +163,17 @@ def box(self, src: Value) -> Value:
         else:
             return src
 
-    def unbox_or_cast(self, src: Value, target_type: RType, line: int) -> Value:
+    def unbox_or_cast(self, src: Value, target_type: RType, line: int, *,
+                      can_borrow: bool = False) -> Value:
         if target_type.is_unboxed:
             return self.add(Unbox(src, target_type, line))
         else:
-            return self.add(Cast(src, target_type, line))
+            if can_borrow:
+                self.keep_alives.append(src)
+            return self.add(Cast(src, target_type, line, borrow=can_borrow))
 
-    def coerce(self, src: Value, target_type: RType, line: int, force: bool = False) -> Value:
+    def coerce(self, src: Value, target_type: RType, line: int, force: bool = False, *,
+               can_borrow: bool = False) -> Value:
         """Generate a coercion/cast from one type to other (only if needed).
 
         For example, int -> object boxes the source int; int -> int emits nothing;
@@ -190,7 +194,7 @@ def coerce(self, src: Value, target_type: RType, line: int, force: bool = False)
             return self.unbox_or_cast(tmp, target_type, line)
         if ((not src.type.is_unboxed and target_type.is_unboxed)
                 or not is_subtype(src.type, target_type)):
-            return self.unbox_or_cast(src, target_type, line)
+            return self.unbox_or_cast(src, target_type, line, can_borrow=can_borrow)
         elif force:
             tmp = Register(target_type)
             self.add(Assign(tmp, src))

mypyc/irbuild/specialize.py

@@ -25,7 +25,7 @@
 )
 from mypyc.ir.rtypes import (
     RType, RTuple, str_rprimitive, list_rprimitive, dict_rprimitive, set_rprimitive,
-    bool_rprimitive, c_int_rprimitive, is_dict_rprimitive
+    bool_rprimitive, c_int_rprimitive, is_dict_rprimitive, is_list_rprimitive
 )
 from mypyc.irbuild.format_str_tokenizer import (
     tokenizer_format_call, join_formatted_strings, convert_format_expr_to_str, FormatOp
@@ -113,14 +113,19 @@ def translate_len(
         builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
     if (len(expr.args) == 1
             and expr.arg_kinds == [ARG_POS]):
-        expr_rtype = builder.node_type(expr.args[0])
+        arg = expr.args[0]
+        expr_rtype = builder.node_type(arg)
         if isinstance(expr_rtype, RTuple):
             # len() of fixed-length tuple can be trivially determined
             # statically, though we still need to evaluate it.
-            builder.accept(expr.args[0])
+            builder.accept(arg)
             return Integer(len(expr_rtype.types))
         else:
-            obj = builder.accept(expr.args[0])
+            if is_list_rprimitive(builder.node_type(arg)):
+                borrow = True
+            else:
+                borrow = False
+            obj = builder.accept(arg, can_borrow=borrow)
             return builder.builtin_len(obj, expr.line)
     return None
 
@@ -429,7 +434,12 @@ def translate_isinstance(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
 
         irs = builder.flatten_classes(expr.args[1])
         if irs is not None:
-            return builder.builder.isinstance_helper(builder.accept(expr.args[0]), irs, expr.line)
+            can_borrow = all(ir.is_ext_class
+                             and not ir.inherits_python
+                             and not ir.allow_interpreted_subclasses
+                             for ir in irs)
+            obj = builder.accept(expr.args[0], can_borrow=can_borrow)
+            return builder.builder.isinstance_helper(obj, irs, expr.line)
     return None
 
 

mypyc/irbuild/statement.py

@@ -58,8 +58,10 @@ 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.
+    # ExpressionStmts do not need to be coerced like other Expressions, so we shouldn't
+    # call builder.accept here.
     stmt.expr.accept(builder.visitor)
+    builder.flush_keep_alives()
 
 
 def transform_return_stmt(builder: IRBuilder, stmt: ReturnStmt) -> None:
@@ -107,6 +109,7 @@ def transform_assignment_stmt(builder: IRBuilder, stmt: AssignmentStmt) -> None:
     for lvalue in lvalues:
         target = builder.get_assignment_target(lvalue)
         builder.assign(target, rvalue_reg, line)
+        builder.flush_keep_alives()
 
 
 def is_simple_lvalue(expr: Expression) -> bool:

mypyc/test-data/exceptions.test

@@ -75,31 +75,28 @@ def f(x):
     r1 :: bit
     r2 :: __main__.A
     r3 :: object
-    r4, r5 :: bit
-    r6 :: int
+    r4 :: bit
+    r5 :: int
 L0:
-    r0 = box(None, 1)
+    r0 = load_address _Py_NoneStruct
     r1 = x == r0
     if r1 goto L1 else goto L2 :: bool
 L1:
     return 2
 L2:
-    inc_ref x
-    r2 = cast(__main__.A, x)
+    r2 = borrow cast(__main__.A, x)
     if is_error(r2) goto L6 (error at f:8) else goto L3
 L3:
-    r3 = box(None, 1)
-    r4 = r2 == r3
-    dec_ref r2
-    r5 = r4 ^ 1
-    if r5 goto L4 else goto L5 :: bool
+    r3 = load_address _Py_NoneStruct
+    r4 = r2 != r3
+    if r4 goto L4 else goto L5 :: bool
 L4:
     return 4
 L5:
     return 6
 L6:
-    r6 = <error> :: int
-    return r6
+    r5 = <error> :: int
+    return r5
 
 [case testListSum]
 from typing import List
@@ -518,4 +515,3 @@ L13:
 L14:
     dec_ref r9
     goto L8
-

mypyc/test-data/irbuild-classes.test

@@ -116,22 +116,22 @@ def Node.length(self):
     self :: __main__.Node
     r0 :: union[__main__.Node, None]
     r1 :: object
-    r2, r3 :: bit
-    r4 :: union[__main__.Node, None]
-    r5 :: __main__.Node
-    r6, r7 :: int
+    r2 :: bit
+    r3 :: union[__main__.Node, None]
+    r4 :: __main__.Node
+    r5, r6 :: int
 L0:
-    r0 = self.next
-    r1 = box(None, 1)
-    r2 = r0 == r1
-    r3 = r2 ^ 1
-    if r3 goto L1 else goto L2 :: bool
+    r0 = borrow self.next
+    r1 = load_address _Py_NoneStruct
+    r2 = r0 != r1
+    keep_alive self
+    if r2 goto L1 else goto L2 :: bool
 L1:
-    r4 = self.next
-    r5 = cast(__main__.Node, r4)
-    r6 = r5.length()
-    r7 = CPyTagged_Add(2, r6)
-    return r7
+    r3 = self.next
+    r4 = cast(__main__.Node, r3)
+    r5 = r4.length()
+    r6 = CPyTagged_Add(2, r5)
+    return r6
 L2:
     return 2