[mypyc] Borrow even more things (#12817)

Borrow operands of tagged integer operations to reduce the number of
incref/decref operations (when it's safe to do so).

Borrow the results in list get item operations, similar to what we've been
doing with get attribute operations.

Author: JukkaL

mypyc/irbuild/ast_helpers.py

@@ -0,0 +1,94 @@
+"""IRBuilder AST transform helpers shared between expressions and statements.
+
+Shared code that is tightly coupled to mypy ASTs can be put here instead of
+making mypyc.irbuild.builder larger.
+"""
+
+from mypy.nodes import (
+    Expression, MemberExpr, Var, IntExpr, FloatExpr, StrExpr, BytesExpr, NameExpr, OpExpr,
+    UnaryExpr, ComparisonExpr, LDEF
+)
+from mypyc.ir.ops import BasicBlock
+from mypyc.ir.rtypes import is_tagged
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.constant_fold import constant_fold_expr
+
+
+def process_conditional(self: IRBuilder, e: Expression, true: BasicBlock,
+                        false: BasicBlock) -> None:
+    if isinstance(e, OpExpr) and e.op in ['and', 'or']:
+        if e.op == 'and':
+            # Short circuit 'and' in a conditional context.
+            new = BasicBlock()
+            process_conditional(self, e.left, new, false)
+            self.activate_block(new)
+            process_conditional(self, e.right, true, false)
+        else:
+            # Short circuit 'or' in a conditional context.
+            new = BasicBlock()
+            process_conditional(self, e.left, true, new)
+            self.activate_block(new)
+            process_conditional(self, e.right, true, false)
+    elif isinstance(e, UnaryExpr) and e.op == 'not':
+        process_conditional(self, e.expr, false, true)
+    else:
+        res = maybe_process_conditional_comparison(self, e, true, false)
+        if res:
+            return
+        # Catch-all for arbitrary expressions.
+        reg = self.accept(e)
+        self.add_bool_branch(reg, true, false)
+
+
+def maybe_process_conditional_comparison(self: IRBuilder,
+                                         e: Expression,
+                                         true: BasicBlock,
+                                         false: BasicBlock) -> bool:
+    """Transform simple tagged integer comparisons in a conditional context.
+
+    Return True if the operation is supported (and was transformed). Otherwise,
+    do nothing and return False.
+
+    Args:
+        e: Arbitrary expression
+        true: Branch target if comparison is true
+        false: Branch target if comparison is false
+    """
+    if not isinstance(e, ComparisonExpr) or len(e.operands) != 2:
+        return False
+    ltype = self.node_type(e.operands[0])
+    rtype = self.node_type(e.operands[1])
+    if not is_tagged(ltype) or not is_tagged(rtype):
+        return False
+    op = e.operators[0]
+    if op not in ('==', '!=', '<', '<=', '>', '>='):
+        return False
+    left_expr = e.operands[0]
+    right_expr = e.operands[1]
+    borrow_left = is_borrow_friendly_expr(self, right_expr)
+    left = self.accept(left_expr, can_borrow=borrow_left)
+    right = self.accept(right_expr, can_borrow=True)
+    # "left op right" for two tagged integers
+    self.builder.compare_tagged_condition(left, right, op, true, false, e.line)
+    return True
+
+
+def is_borrow_friendly_expr(self: 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 immortal and can always be borrowed
+        return True
+    if (isinstance(expr, (UnaryExpr, OpExpr, NameExpr, MemberExpr)) and
+            constant_fold_expr(self, 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 self.is_native_attr_ref(expr):
+        return True
+    return False

mypyc/irbuild/builder.py

@@ -14,15 +14,15 @@
 
 from mypyc.irbuild.prepare import RegisterImplInfo
 from typing import Callable, Dict, List, Tuple, Optional, Union, Sequence, Set, Any, Iterator
-from typing_extensions import overload
+from typing_extensions import overload, Final
 from mypy.backports import OrderedDict
 
 from mypy.build import Graph
 from mypy.nodes import (
     MypyFile, SymbolNode, Statement, OpExpr, IntExpr, NameExpr, LDEF, Var, UnaryExpr,
     CallExpr, IndexExpr, Expression, MemberExpr, RefExpr, Lvalue, TupleExpr,
-    TypeInfo, Decorator, OverloadedFuncDef, StarExpr, ComparisonExpr, GDEF,
-    ArgKind, ARG_POS, ARG_NAMED, FuncDef,
+    TypeInfo, Decorator, OverloadedFuncDef, StarExpr,
+    GDEF, ArgKind, ARG_POS, ARG_NAMED, FuncDef,
 )
 from mypy.types import (
     Type, Instance, TupleType, UninhabitedType, get_proper_type
@@ -40,7 +40,7 @@
 from mypyc.ir.rtypes import (
     RType, RTuple, RInstance, c_int_rprimitive, int_rprimitive, dict_rprimitive,
     none_rprimitive, is_none_rprimitive, object_rprimitive, is_object_rprimitive,
-    str_rprimitive, is_tagged, is_list_rprimitive, is_tuple_rprimitive, c_pyssize_t_rprimitive
+    str_rprimitive, is_list_rprimitive, is_tuple_rprimitive, c_pyssize_t_rprimitive
 )
 from mypyc.ir.func_ir import FuncIR, INVALID_FUNC_DEF, RuntimeArg, FuncSignature, FuncDecl
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
@@ -67,6 +67,11 @@
 from mypyc.irbuild.util import is_constant
 
 
+# These int binary operations can borrow their operands safely, since the
+# primitives take this into consideration.
+int_borrow_friendly_op: Final = {'+', '-', '==', '!=', '<', '<=', '>', '>='}
+
+
 class IRVisitor(ExpressionVisitor[Value], StatementVisitor[None]):
     pass
 
@@ -287,7 +292,7 @@ def gen_method_call(self,
                         arg_kinds: Optional[List[ArgKind]] = None,
                         arg_names: Optional[List[Optional[str]]] = None) -> Value:
         return self.builder.gen_method_call(
-            base, name, arg_values, result_type, line, arg_kinds, arg_names
+            base, name, arg_values, result_type, line, arg_kinds, arg_names, self.can_borrow
         )
 
     def load_module(self, name: str) -> Value:
@@ -515,7 +520,7 @@ def get_assignment_target(self, lvalue: Lvalue,
             # Attribute assignment x.y = e
             can_borrow = self.is_native_attr_ref(lvalue)
             obj = self.accept(lvalue.expr, can_borrow=can_borrow)
-            return AssignmentTargetAttr(obj, lvalue.name)
+            return AssignmentTargetAttr(obj, lvalue.name, can_borrow=can_borrow)
         elif isinstance(lvalue, TupleExpr):
             # Multiple assignment a, ..., b = e
             star_idx: Optional[int] = None
@@ -535,7 +540,10 @@ def get_assignment_target(self, lvalue: Lvalue,
 
         assert False, 'Unsupported lvalue: %r' % lvalue
 
-    def read(self, target: Union[Value, AssignmentTarget], line: int = -1) -> Value:
+    def read(self,
+             target: Union[Value, AssignmentTarget],
+             line: int = -1,
+             can_borrow: bool = False) -> Value:
         if isinstance(target, Value):
             return target
         if isinstance(target, AssignmentTargetRegister):
@@ -548,7 +556,8 @@ def read(self, target: Union[Value, AssignmentTarget], line: int = -1) -> Value:
             assert False, target.base.type
         if isinstance(target, AssignmentTargetAttr):
             if isinstance(target.obj.type, RInstance) and target.obj.type.class_ir.is_ext_class:
-                return self.add(GetAttr(target.obj, target.attr, line))
+                borrow = can_borrow and target.can_borrow
+                return self.add(GetAttr(target.obj, target.attr, line, borrow=borrow))
             else:
                 return self.py_get_attr(target.obj, target.attr, line)
 
@@ -915,61 +924,6 @@ def shortcircuit_expr(self, expr: OpExpr) -> Value:
             expr.line
         )
 
-    # Conditional expressions
-
-    def process_conditional(self, e: Expression, true: BasicBlock, false: BasicBlock) -> None:
-        if isinstance(e, OpExpr) and e.op in ['and', 'or']:
-            if e.op == 'and':
-                # Short circuit 'and' in a conditional context.
-                new = BasicBlock()
-                self.process_conditional(e.left, new, false)
-                self.activate_block(new)
-                self.process_conditional(e.right, true, false)
-            else:
-                # Short circuit 'or' in a conditional context.
-                new = BasicBlock()
-                self.process_conditional(e.left, true, new)
-                self.activate_block(new)
-                self.process_conditional(e.right, true, false)
-        elif isinstance(e, UnaryExpr) and e.op == 'not':
-            self.process_conditional(e.expr, false, true)
-        else:
-            res = self.maybe_process_conditional_comparison(e, true, false)
-            if res:
-                return
-            # Catch-all for arbitrary expressions.
-            reg = self.accept(e)
-            self.add_bool_branch(reg, true, false)
-
-    def maybe_process_conditional_comparison(self,
-                                             e: Expression,
-                                             true: BasicBlock,
-                                             false: BasicBlock) -> bool:
-        """Transform simple tagged integer comparisons in a conditional context.
-
-        Return True if the operation is supported (and was transformed). Otherwise,
-        do nothing and return False.
-
-        Args:
-            e: Arbitrary expression
-            true: Branch target if comparison is true
-            false: Branch target if comparison is false
-        """
-        if not isinstance(e, ComparisonExpr) or len(e.operands) != 2:
-            return False
-        ltype = self.node_type(e.operands[0])
-        rtype = self.node_type(e.operands[1])
-        if not is_tagged(ltype) or not is_tagged(rtype):
-            return False
-        op = e.operators[0]
-        if op not in ('==', '!=', '<', '<=', '>', '>='):
-            return False
-        left = self.accept(e.operands[0])
-        right = self.accept(e.operands[1])
-        # "left op right" for two tagged integers
-        self.builder.compare_tagged_condition(left, right, op, true, false, e.line)
-        return True
-
     # Basic helpers
 
     def flatten_classes(self, arg: Union[RefExpr, TupleExpr]) -> Optional[List[ClassIR]]:

mypyc/irbuild/expression.py

@@ -40,12 +40,13 @@
 from mypyc.primitives.str_ops import str_slice_op
 from mypyc.primitives.int_ops import int_comparison_op_mapping
 from mypyc.irbuild.specialize import apply_function_specialization, apply_method_specialization
-from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
 from mypyc.irbuild.for_helpers import (
     translate_list_comprehension, translate_set_comprehension,
     comprehension_helper
 )
 from mypyc.irbuild.constant_fold import constant_fold_expr
+from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
 
 
 # Name and attribute references
@@ -404,6 +405,15 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
     if folded:
         return folded
 
+    # Special case some int ops to allow borrowing operands.
+    if (is_int_rprimitive(builder.node_type(expr.left))
+            and is_int_rprimitive(builder.node_type(expr.right))):
+        if expr.op in int_borrow_friendly_op:
+            borrow_left = is_borrow_friendly_expr(builder, expr.right)
+            left = builder.accept(expr.left, can_borrow=borrow_left)
+            right = builder.accept(expr.right, can_borrow=True)
+            return builder.binary_op(left, right, expr.op, expr.line)
+
     return builder.binary_op(
         builder.accept(expr.left), builder.accept(expr.right), expr.op, expr.line
     )
@@ -430,26 +440,6 @@ def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
         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.
 
@@ -504,7 +494,7 @@ def try_gen_slice_op(builder: IRBuilder, base: Value, index: SliceExpr) -> Optio
 def transform_conditional_expr(builder: IRBuilder, expr: ConditionalExpr) -> Value:
     if_body, else_body, next_block = BasicBlock(), BasicBlock(), BasicBlock()
 
-    builder.process_conditional(expr.cond, if_body, else_body)
+    process_conditional(builder, expr.cond, if_body, else_body)
     expr_type = builder.node_type(expr)
     # Having actual Phi nodes would be really nice here!
     target = Register(expr_type)
@@ -577,11 +567,22 @@ 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')
+    if len(e.operators) == 1:
+        # Special some common simple cases
+        if first_op in ('is', 'is not'):
+            right_expr = e.operands[1]
+            if isinstance(right_expr, NameExpr) and right_expr.fullname == 'builtins.None':
+                # Special case 'is None' / 'is not None'.
+                return translate_is_none(builder, e.operands[0], negated=first_op != 'is')
+        left_expr = e.operands[0]
+        if is_int_rprimitive(builder.node_type(left_expr)):
+            right_expr = e.operands[1]
+            if is_int_rprimitive(builder.node_type(right_expr)):
+                if first_op in int_borrow_friendly_op:
+                    borrow_left = is_borrow_friendly_expr(builder, right_expr)
+                    left = builder.accept(left_expr, can_borrow=borrow_left)
+                    right = builder.accept(right_expr, can_borrow=True)
+                    return builder.compare_tagged(left, right, first_op, e.line)
 
     # TODO: Don't produce an expression when used in conditional context
     # All of the trickiness here is due to support for chained conditionals