[mypyc] Borrow references during chained attribute access

mypyc/codegen/emitfunc.py

@@ -333,7 +333,7 @@ def visit_get_attr(self, op: GetAttr) -> None:
                         'PyErr_SetString({}, "attribute {} of {} undefined");'.format(
                             exc_class, repr(op.attr), repr(cl.name)))
 
-            if attr_rtype.is_refcounted:
+            if attr_rtype.is_refcounted and not op.is_borrowed:
                 if not merged_branch and not always_defined:
                     self.emitter.emit_line('} else {')
                 self.emitter.emit_inc_ref(dest, attr_rtype)

mypyc/ir/ops.py

@@ -599,13 +599,14 @@ class GetAttr(RegisterOp):
 
     error_kind = ERR_MAGIC
 
-    def __init__(self, obj: Value, attr: str, line: int) -> None:
+    def __init__(self, obj: Value, attr: str, line: int, *, borrow: bool = False) -> None:
         super().__init__(line)
         self.obj = obj
         self.attr = attr
         assert isinstance(obj.type, RInstance), 'Attribute access not supported: %s' % obj.type
         self.class_type = obj.type
         self.type = obj.type.attr_type(attr)
+        self.is_borrowed = borrow
 
     def sources(self) -> List[Value]:
         return [self.obj]

mypyc/ir/pprint.py

@@ -77,7 +77,11 @@ 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 = %r.%s', op, op.obj, op.attr)
+        if op.is_borrowed:
+            borrow = 'borrow '
+        else:
+            borrow = ''
+        return self.format('%r = %s%r.%s', op, borrow, op.obj, op.attr)
 
     def visit_set_attr(self, op: SetAttr) -> str:
         if op.is_init:

mypyc/irbuild/builder.py

@@ -141,6 +141,8 @@ def __init__(self,
         # can also do quick lookups.
         self.imports: OrderedDict[str, None] = OrderedDict()
 
+        self.can_borrow = False
+
     # High-level control
 
     def set_module(self, module_name: str, module_path: str) -> None:
@@ -152,15 +154,23 @@ def set_module(self, module_name: str, module_path: str) -> None:
         self.module_path = module_path
 
     @overload
-    def accept(self, node: Expression) -> Value: ...
+    def accept(self, node: Expression, *, can_borrow: bool = False) -> Value: ...
 
     @overload
     def accept(self, node: Statement) -> None: ...
 
-    def accept(self, node: Union[Statement, Expression]) -> Optional[Value]:
-        """Transform an expression or a statement."""
+    def accept(self, node: Union[Statement, Expression], *,
+               can_borrow: bool = False) -> Optional[Value]:
+        """Transform an expression or a statement.
+
+        If can_borrow is true, prefer to generate a borrowed reference.
+        Borrowed references are faster since they don't require reference count
+        manipulation, but they are only safe to use in specific contexts.
+        """
         with self.catch_errors(node.line):
             if isinstance(node, Expression):
+                old_can_borrow = self.can_borrow
+                self.can_borrow = can_borrow
                 try:
                     res = node.accept(self.visitor)
                     res = self.coerce(res, self.node_type(node), node.line)
@@ -170,6 +180,9 @@ def accept(self, node: Union[Statement, Expression]) -> Optional[Value]:
                 # from causing more downstream trouble.
                 except UnsupportedException:
                     res = Register(self.node_type(node))
+                self.can_borrow = old_can_borrow
+                if not can_borrow:
+                    self.builder.flush_keep_alives()
                 return res
             else:
                 try:

mypyc/irbuild/expression.py

@@ -21,7 +21,7 @@
     Value, Register, TupleGet, TupleSet, BasicBlock, Assign, LoadAddress, RaiseStandardError
 )
 from mypyc.ir.rtypes import (
-    RTuple, object_rprimitive, is_none_rprimitive, int_rprimitive, is_int_rprimitive
+    RTuple, RInstance, object_rprimitive, is_none_rprimitive, int_rprimitive, is_int_rprimitive
 )
 from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD
 from mypyc.irbuild.format_str_tokenizer import (
@@ -130,8 +130,19 @@ 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 = builder.accept(expr.expr)
+    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
+    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.
     typ = get_proper_type(builder.types.get(expr.expr))
     if isinstance(typ, TupleType) and typ.partial_fallback.type.is_named_tuple:
@@ -142,7 +153,8 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
 
     check_instance_attribute_access_through_class(builder, expr, typ)
 
-    return builder.builder.get_attr(obj, expr.name, rtype, expr.line)
+    borrow = can_borrow and builder.can_borrow
+    return builder.builder.get_attr(obj, expr.name, rtype, expr.line, borrow=borrow)
 
 
 def check_instance_attribute_access_through_class(builder: IRBuilder,

mypyc/irbuild/ll_builder.py

@@ -105,6 +105,9 @@ def __init__(
         self.blocks: List[BasicBlock] = []
         # Stack of except handler entry blocks
         self.error_handlers: List[Optional[BasicBlock]] = [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] = []
 
     # Basic operations
 
@@ -145,6 +148,11 @@ def self(self) -> Register:
         """
         return self.args[0]
 
+    def flush_keep_alives(self) -> None:
+        if self.keep_alives:
+            self.add(KeepAlive(self.keep_alives[:]))
+            self.keep_alives = []
+
     # Type conversions
 
     def box(self, src: Value) -> Value:
@@ -219,11 +227,14 @@ 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) -> Value:
+    def get_attr(self, obj: Value, attr: str, result_type: RType, line: int, *,
+                 borrow: bool = False) -> Value:
         """Get a native or Python attribute of an object."""
         if (isinstance(obj.type, RInstance) and obj.type.class_ir.is_ext_class
                 and obj.type.class_ir.has_attr(attr)):
-            return self.add(GetAttr(obj, attr, line))
+            if borrow:
+                self.keep_alives.append(obj)
+            return self.add(GetAttr(obj, attr, line, borrow=borrow))
         elif isinstance(obj.type, RUnion):
             return self.union_get_attr(obj, obj.type, attr, result_type, line)
         else:

mypyc/test-data/irbuild-classes.test

@@ -1240,3 +1240,51 @@ L0:
     r0 = ''
     __mypyc_self__.s = r0
     return 1
+
+[case testBorrowAttribute]
+def f(d: D) -> int:
+    return d.c.x
+
+class C:
+    x: int
+class D:
+    c: C
+[out]
+def f(d):
+    d :: __main__.D
+    r0 :: __main__.C
+    r1 :: int
+L0:
+    r0 = borrow d.c
+    r1 = r0.x
+    keep_alive d
+    return r1
+
+[case testNoBorrowOverPropertyAccess]
+class C:
+    d: D
+class D:
+    @property
+    def e(self) -> E:
+        return E()
+class E:
+    x: int
+def f(c: C) -> int:
+    return c.d.e.x
+[out]
+def D.e(self):
+    self :: __main__.D
+    r0 :: __main__.E
+L0:
+    r0 = E()
+    return r0
+def f(c):
+    c :: __main__.C
+    r0 :: __main__.D
+    r1 :: __main__.E
+    r2 :: int
+L0:
+    r0 = c.d
+    r1 = r0.e
+    r2 = r1.x
+    return r2

mypyc/test-data/refcount.test

@@ -917,3 +917,36 @@ L0:
     r5 = unbox(int, r4)
     dec_ref r4
     return r5
+
+[case testBorrowAttribute]
+def g() -> int:
+    d = D()
+    return d.c.x
+
+def f(d: D) -> int:
+    return d.c.x
+
+class C:
+    x: int
+class D:
+    c: C
+[out]
+def g():
+    r0, d :: __main__.D
+    r1 :: __main__.C
+    r2 :: int
+L0:
+    r0 = D()
+    d = r0
+    r1 = borrow d.c
+    r2 = r1.x
+    dec_ref d
+    return r2
+def f(d):
+    d :: __main__.D
+    r0 :: __main__.C
+    r1 :: int
+L0:
+    r0 = borrow d.c
+    r1 = r0.x
+    return r1