Make Builtin.addressForBorrow work with addressable parameters.

Author: jckarter

lib/SILGen/SILGenApply.cpp

@@ -3410,6 +3410,34 @@ Expr *ArgumentSource::findStorageReferenceExprForBorrow() && {
   return lvExpr;
 }
 
+ManagedValue
+SILGenFunction::tryEmitAddressableParameterAsAddress(ArgumentSource &&arg,
+                                                     ValueOwnership ownership) {
+  // If the function takes an addressable parameter, and its argument is
+  // a reference to an addressable declaration with compatible ownership,
+  // forward the address along in-place.
+  if (arg.isExpr()) {
+    auto origExpr = std::move(arg).asKnownExpr();
+    auto expr = origExpr;
+    
+    if (auto le = dyn_cast<LoadExpr>(expr)) {
+      expr = le->getSubExpr();
+    }
+    if (auto dre = dyn_cast<DeclRefExpr>(expr)) {
+      if (auto param = dyn_cast<ParamDecl>(dre->getDecl())) {
+        if (VarLocs.count(param)
+            && VarLocs[param].addressable
+            && param->getValueOwnership() == ownership) {
+          auto addr = VarLocs[param].value;
+          return ManagedValue::forBorrowedAddressRValue(addr);
+        }
+      }
+    }
+    arg = ArgumentSource(origExpr);
+  }
+  return ManagedValue();
+}
+
 namespace {
 
 class ArgEmitter {
@@ -3515,29 +3543,11 @@ class ArgEmitter {
       // If the function takes an addressable parameter, and its argument is
       // a reference to an addressable declaration with compatible ownership,
       // forward the address along in-place.
-      if (arg.isExpr()) {
-        arg.dump();
-        auto origExpr = std::move(arg).asKnownExpr();
-        auto expr = origExpr;
-        
-        if (auto le = dyn_cast<LoadExpr>(expr)) {
-          expr = le->getSubExpr();
-        }
-        if (auto dre = dyn_cast<DeclRefExpr>(expr)) {
-          if (auto param = dyn_cast<ParamDecl>(dre->getDecl())) {
-            if (SGF.VarLocs.count(param)
-              && SGF.VarLocs[param].addressable
-              && param->getValueOwnership() == origParam->getValueOwnership()) {
-              auto addr = SGF.VarLocs[param].value;
-              claimNextParameters(1);
-              Args.push_back(ManagedValue::forBorrowedAddressRValue(addr));
-              return;
-            }
-          }
-        }
-        llvm::errs() << "couldn't lower as addressable\n";
-        arg = ArgumentSource(origExpr);
-        arg.dump();
+      if (auto addr = SGF.tryEmitAddressableParameterAsAddress(std::move(arg),
+                                             origParam->getValueOwnership())) {
+        claimNextParameters(1);
+        Args.push_back(addr);
+        return;
       }
     }
 

lib/SILGen/SILGenBuiltin.cpp

@@ -492,27 +492,38 @@ static ManagedValue emitBuiltinAddressOfBorrowBuiltins(SILGenFunction &SGF,
   auto argument = (*argsOrError)[0];
 
   SILValue addr;
-  // Try to borrow the argument at +0. We only support if it's
-  // naturally emitted borrowed in memory.
-  auto borrow = SGF.emitRValue(argument, SGFContext::AllowGuaranteedPlusZero)
-     .getAsSingleValue(SGF, argument);
-  if (!SGF.F.getConventions().useLoweredAddresses()) {
-    auto &context = SGF.getASTContext();
-    auto identifier =
-        stackProtected
-            ? context.getIdentifier("addressOfBorrowOpaque")
-            : context.getIdentifier("unprotectedAddressOfBorrowOpaque");
-    auto builtin = SGF.B.createBuiltin(loc, identifier, rawPointerType,
-                                       substitutions, {borrow.getValue()});
-    return ManagedValue::forObjectRValueWithoutOwnership(builtin);
-  }
+  // Try to borrow the argument at +0 indirect.
+  // If the argument is a reference to a borrowed addressable parameter, then
+  // use that parameter's stable address.
+  if (auto addressableAddr = SGF.tryEmitAddressableParameterAsAddress(
+                                                      ArgumentSource(argument),
+                                                      ValueOwnership::Shared)) {
+    addr = addressableAddr.getValue();
+  } else {
+    // We otherwise only support the builtin applied to values that
+    // are naturally emitted borrowed in memory. (But it would probably be good
+    // to phase this out since it's not really well-defined how long
+    // the resulting pointer is good for without something like addressability.)
+    auto borrow = SGF.emitRValue(argument, SGFContext::AllowGuaranteedPlusZero)
+       .getAsSingleValue(SGF, argument);
+    if (!SGF.F.getConventions().useLoweredAddresses()) {
+      auto &context = SGF.getASTContext();
+      auto identifier =
+          stackProtected
+              ? context.getIdentifier("addressOfBorrowOpaque")
+              : context.getIdentifier("unprotectedAddressOfBorrowOpaque");
+      auto builtin = SGF.B.createBuiltin(loc, identifier, rawPointerType,
+                                         substitutions, {borrow.getValue()});
+      return ManagedValue::forObjectRValueWithoutOwnership(builtin);
+    }
 
-  if (!borrow.isPlusZero() || !borrow.getType().isAddress()) {
-    SGF.SGM.diagnose(argument->getLoc(), diag::non_borrowed_indirect_addressof);
-    return SGF.emitUndef(rawPointerType);
-  }
+    if (!borrow.isPlusZero() || !borrow.getType().isAddress()) {
+      SGF.SGM.diagnose(argument->getLoc(), diag::non_borrowed_indirect_addressof);
+      return SGF.emitUndef(rawPointerType);
+    }
 
-  addr = borrow.getValue();
+    addr = borrow.getValue();
+  }
 
   // Take the address argument and cast it to RawPointer.
   SILValue result = SGF.B.createAddressToPointer(loc, addr, rawPointerType,

lib/SILGen/SILGenFunction.h

@@ -1424,6 +1424,16 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   ManagedValue manageBufferForExprResult(SILValue buffer,
                                          const TypeLowering &bufferTL,
                                          SGFContext C);
+                                         
+  /// Tries to emit an argument referring to an addressable parameter as the
+  /// stable address of the parameter.
+  ///
+  /// Returns a null ManagedValue if the argument is not a parameter reference,
+  /// the referenced parameter is not addressable, or the requested
+  /// \c ownership is not compatible with the parameter's ownership. \c arg
+  /// is consumed only if the operation succeeds.
+  ManagedValue tryEmitAddressableParameterAsAddress(ArgumentSource &&arg,
+                                                    ValueOwnership ownership);
 
   //===--------------------------------------------------------------------===//
   // Type conversions for expr emission and thunks

test/SILOptimizer/addressable_dependencies.swift

@@ -0,0 +1,38 @@
+// RUN: %target-swift-frontend -emit-sil -enable-experimental-feature BuiltinModule -enable-experimental-feature LifetimeDependence -enable-experimental-feature AddressableTypes %s | %FileCheck %s
+
+// REQUIRES: swift_feature_BuiltinModule
+// REQUIRES: swift_feature_AddressableTypes
+// REQUIRES: swift_feature_LifetimeDependence
+
+import Builtin
+
+@_addressableForDependencies
+struct Node {
+    var id: String
+
+    var ref: NodeRef {
+      // CHECK-LABEL: sil {{.*}}@${{.*}}4NodeV3ref{{.*}}Vvg :
+      // CHECK-SAME:    (@in_guaranteed Node) ->
+      @lifetime(borrow self)
+      borrowing get {
+        // CHECK: bb0(%0 : @noImplicitCopy $*Node):
+        // CHECK: [[REF:%.*]] = apply {{.*}}(%0,
+        // CHECK: mark_dependence [nonescaping] [[REF]] on %0
+        return NodeRef(node: self)
+      }
+    }
+}
+
+struct NodeRef: ~Escapable {
+    private var parent: UnsafePointer<Node>
+
+    // CHECK-LABEL: sil {{.*}}@${{.*}}7NodeRefV4node{{.*}}fC :
+    // CHECK-SAME:    (@in_guaranteed Node,
+    @lifetime(borrow node)
+    init(node: borrowing Node) {
+        // CHECK: bb0(%0 : @noImplicitCopy $*Node,
+        // CHECK:   [[RAW_PTR:%.*]] = address_to_pointer {{.*}}%0
+        // CHECK:   struct $UnsafePointer<Node> ([[RAW_PTR]])
+        self.parent = UnsafePointer(Builtin.addressOfBorrow(node))
+    }
+}