Add a tiny BorrowedAddress utility.

Determines whether an address might be inside a borrowed scope. If so,
then any address substitution needs to find that scope boundary to
avoid violating its basic guarantee that all uses are within scope.

Author: atrick

include/swift/SIL/InstructionUtils.h

@@ -27,10 +27,6 @@ SILValue getUnderlyingObject(SILValue V);
 
 SILValue getUnderlyingObjectStopAtMarkDependence(SILValue V);
 
-/// Given an address look through address to address projections and indexing
-/// insts.
-SILValue getUnderlyingObjectStoppingAtObjectToAddrProjections(SILValue v);
-
 SILValue stripSinglePredecessorArgs(SILValue V);
 
 /// Return the underlying SILValue after stripping off all casts from the

include/swift/SIL/OwnershipUtils.h

@@ -727,6 +727,8 @@ struct InteriorPointerOperand {
     return kind != InteriorPointerOperandKind::Invalid && operand;
   }
 
+  SILInstruction *getUser() const { return operand->getUser(); }
+
   /// If \p op has a user that is an interior pointer, return a valid
   /// value. Otherwise, return None.
   static InteriorPointerOperand get(Operand *op) {
@@ -835,6 +837,25 @@ struct InteriorPointerOperand {
       : operand(op), kind(kind) {}
 };
 
+/// Utility to check if an address may originate from a borrowed value. If so,
+/// then uses of the address cannot be replaced without ensuring that they are
+/// also within the same scope.
+///
+/// If mayBeBorrowed is false, then there is no enclosing borrow scope and
+/// interiorPointerOp is irrelevant.
+///
+/// If mayBeBorrowed is true, then interiorPointerOp refers to the operand that
+/// converts a non-address value into the address from which the contructor's
+/// address is derived. If the best-effort to find an InteriorPointerOperand
+/// fails, then interiorPointerOp remains invalid, and clients must be
+/// conservative.
+struct BorrowedAddress {
+  bool mayBeBorrowed = true;
+  InteriorPointerOperand interiorPointerOp;
+
+  BorrowedAddress(SILValue address);
+};
+
 class OwnedValueIntroducerKind {
 public:
   enum Kind : uint8_t {

lib/SIL/Utils/InstructionUtils.cpp

@@ -62,20 +62,6 @@ SILValue swift::getUnderlyingObject(SILValue v) {
   }
 }
 
-SILValue
-swift::getUnderlyingObjectStoppingAtObjectToAddrProjections(SILValue v) {
-  if (!v->getType().isAddress())
-    return SILValue();
-
-  while (true) {
-    auto v2 = lookThroughAddressToAddressProjections(v);
-    v2 = stripIndexingInsts(v2);
-    if (v2 == v)
-      return v2;
-    v = v2;
-  }
-}
-
 SILValue swift::getUnderlyingObjectStopAtMarkDependence(SILValue v) {
   while (true) {
     SILValue v2 = stripCastsWithoutMarkDependence(v);

lib/SIL/Utils/OwnershipUtils.cpp

@@ -16,6 +16,7 @@
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/LinearLifetimeChecker.h"
 #include "swift/SIL/Projection.h"
+#include "swift/SIL/MemAccessUtils.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILInstruction.h"
 
@@ -787,6 +788,46 @@ bool InteriorPointerOperand::findTransitiveUsesForAddress(
   return foundError;
 }
 
+// Determine whether \p address may be in a borrow scope and if so record the
+// InteriorPointerOperand that produces the address from a guaranteed value.
+BorrowedAddress::BorrowedAddress(SILValue address) {
+  assert(address->getType().isAddress());
+
+  auto storageWithBase = AccessedStorageWithBase::compute(address);
+  switch (storageWithBase.storage.getKind()) {
+  case AccessedStorage::Argument:
+  case AccessedStorage::Stack:
+  case AccessedStorage::Global:
+  // Unidentified storage is from an "escaped pointer", so it need not be
+  // restricted to a borrow scope.
+  case AccessedStorage::Unidentified:
+    this->mayBeBorrowed = false;
+    return;
+  case AccessedStorage::Box:
+  case AccessedStorage::Yield:
+  case AccessedStorage::Class:
+  case AccessedStorage::Tail:
+    // The base object might be within a borrow scope.
+    break;
+  case AccessedStorage::Nested:
+    llvm_unreachable("unexpected storage");
+  };
+  auto base = storageWithBase.base;
+  if (!base)
+    return; // bail on unexpected patterns
+
+  auto intPtrOp = InteriorPointerOperand::inferFromResult(base);
+  if (!intPtrOp)
+    return;
+
+  SILValue nonAddressValue = intPtrOp.operand->get();
+  if (nonAddressValue->getOwnershipKind() != OwnershipKind::Guaranteed) {
+    this->mayBeBorrowed = false;
+    return;
+  }
+  this->interiorPointerOp = intPtrOp;
+}
+
 //===----------------------------------------------------------------------===//
 //                          Owned Value Introducers
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -903,41 +903,31 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
   //
   // NOTE: We also need to handle this here since a pointer_to_address is not a
   // valid base value for an access path since it doesn't refer to any storage.
-  {
-    auto baseProj =
-        getUnderlyingObjectStoppingAtObjectToAddrProjections(newValue);
-    if (isa<PointerToAddressInst>(baseProj)) {
-      return;
-    }
-  }
+  BorrowedAddress borrowedAddress(newValue);
+  if (!borrowedAddress.mayBeBorrowed)
+    return;
 
-  auto accessPathWithBase = AccessPathWithBase::compute(newValue);
-  if (!accessPathWithBase.base) {
+  if (!borrowedAddress.interiorPointerOp) {
     // Invalidate!
     ctx = nullptr;
     return;
   }
 
-  auto &intPtr = ctx->extraAddressFixupInfo.intPtrOp;
-  intPtr = InteriorPointerOperand::inferFromResult(accessPathWithBase.base);
-  if (!intPtr) {
-    // We can optimize! Do not invalidate!
-    return;
-  }
-
-  auto borrowedValue = intPtr.getSingleBaseValue();
+  ctx->extraAddressFixupInfo.intPtrOp = borrowedAddress.interiorPointerOp;
+  auto borrowedValue = borrowedAddress.interiorPointerOp.getSingleBaseValue();
   if (!borrowedValue) {
     // Invalidate!
     ctx = nullptr;
     return;
   }
 
-  // This cloner check must match the later cloner invocation in
-  // replaceAddressUses()
-  auto *intPtrUser = cast<SingleValueInstruction>(intPtr->getUser());
+  auto *intPtrInst =
+    cast<SingleValueInstruction>(borrowedAddress.interiorPointerOp.getUser());
   auto checkBase = [&](SILValue srcAddr) {
-    return (srcAddr == intPtrUser) ? SILValue(intPtrUser) : SILValue();
+    return (srcAddr == intPtrInst) ? SILValue(intPtrInst) : SILValue();
   };
+  // This cloner check must match the later cloner invocation in
+  // replaceAddressUses()
   if (!canCloneUseDefChain(newValue, checkBase)) {
     ctx = nullptr;
     return;