OSSA ownership optimization RAUW utility fixes.

Verify that the OwnershipRAUWUtility always preserves the original
borrow scope by exhaustively switching over OperandOwnership.

And related cleanup.

include/swift/SIL/OwnershipUtils.h

@@ -76,12 +76,10 @@ inline bool isForwardingConsume(SILValue value) {
 }
 
 class ForwardingOperand {
-  Operand *use;
-
-  ForwardingOperand(Operand *use) : use(use) {}
+  Operand *use = nullptr;
 
 public:
-  static ForwardingOperand get(Operand *use);
+  explicit ForwardingOperand(Operand *use);
 
   OwnershipConstraint getOwnershipConstraint() const {
     // We use a force unwrap since a ForwardingOperand should always have an
@@ -665,23 +663,24 @@ struct InteriorPointerOperand {
     llvm_unreachable("Covered switch isn't covered?!");
   }
 
-  /// Compute the list of implicit uses that this interior pointer operand puts
-  /// on its parent guaranted value.
+  /// Transitively compute the list of uses that this interior pointer operand
+  /// puts on its parent guaranted value.
   ///
   /// Example: Uses of a ref_element_addr can not occur outside of the lifetime
   /// of the instruction's operand. The uses of that address act as liveness
   /// requirements to ensure that the underlying class is alive at all use
   /// points.
-  bool getImplicitUses(SmallVectorImpl<Operand *> &foundUses,
+  bool findTransitiveUses(SmallVectorImpl<Operand *> &foundUses,
                           std::function<void(Operand *)> *onError = nullptr) {
-    return getImplicitUsesForAddress(getProjectedAddress(), foundUses, onError);
+    return findTransitiveUsesForAddress(getProjectedAddress(), foundUses,
+                                        onError);
   }
 
   /// The algorithm that is used to determine what the verifier will consider to
-  /// be implicit uses of the given address. Used to implement \see
-  /// getImplicitUses.
+  /// be transitive uses of the given address. Used to implement \see
+  /// findTransitiveUses.
   static bool
-  getImplicitUsesForAddress(SILValue address,
+  findTransitiveUsesForAddress(SILValue address,
                                SmallVectorImpl<Operand *> &foundUses,
                                std::function<void(Operand *)> *onError = nullptr);
 

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -38,6 +38,9 @@ struct OwnershipFixupContext {
   DeadEndBlocks &deBlocks;
   JointPostDominanceSetComputer &jointPostDomSetComputer;
 
+  SmallVector<Operand *, 8> transitiveBorrowedUses;
+  SmallVector<BorrowingOperand, 8> recursiveReborrows;
+
   /// Extra state initialized by OwnershipRAUWFixupHelper::get() that we use
   /// when RAUWing addresses. This ensures we do not need to recompute this
   /// state when we perform the actual RAUW.
@@ -67,6 +70,8 @@ struct OwnershipFixupContext {
 
   void clear() {
     jointPostDomSetComputer.clear();
+    transitiveBorrowedUses.clear();
+    recursiveReborrows.clear();
     extraAddressFixupInfo.allAddressUsesFromOldValue.clear();
     extraAddressFixupInfo.intPtrOp = InteriorPointerOperand();
   }

lib/SIL/Utils/OwnershipUtils.cpp

@@ -520,7 +520,7 @@ bool BorrowedValue::visitInteriorPointerOperands(
 //                           InteriorPointerOperand
 //===----------------------------------------------------------------------===//
 
-bool InteriorPointerOperand::getImplicitUsesForAddress(
+bool InteriorPointerOperand::findTransitiveUsesForAddress(
     SILValue projectedAddress, SmallVectorImpl<Operand *> &foundUses,
     std::function<void(Operand *)> *onError) {
   SmallVector<Operand *, 8> worklist(projectedAddress->getUses());
@@ -882,12 +882,12 @@ OwnedValueIntroducer swift::getSingleOwnedValueIntroducer(SILValue inputValue) {
 //                             Forwarding Operand
 //===----------------------------------------------------------------------===//
 
-ForwardingOperand ForwardingOperand::get(Operand *use) {
+ForwardingOperand::ForwardingOperand(Operand *use) {
   if (use->isTypeDependent())
-    return nullptr;
+    return;
 
   if (!OwnershipForwardingMixin::isa(use->getUser())) {
-    return nullptr;
+    return;
   }
 #ifndef NDEBUG
   switch (use->getOperandOwnership()) {
@@ -909,7 +909,7 @@ ForwardingOperand ForwardingOperand::get(Operand *use) {
     llvm_unreachable("this isn't the operand being forwarding!");
   }
 #endif
-  return {use};
+  this->use = use;
 }
 
 ValueOwnershipKind ForwardingOperand::getOwnershipKind() const {
@@ -1008,7 +1008,7 @@ void ForwardingOperand::setOwnershipKind(ValueOwnershipKind newKind) const {
     return;
   }
 
-  llvm_unreachable("Out of sync with ForwardingOperand::get?!");
+  llvm_unreachable("Out of sync with OperandOwnership");
 }
 
 void ForwardingOperand::replaceOwnershipKind(ValueOwnershipKind oldKind,
@@ -1075,7 +1075,7 @@ void ForwardingOperand::replaceOwnershipKind(ValueOwnershipKind oldKind,
     return;
   }
 
-  llvm_unreachable("Missing Case! Out of sync with ForwardingOperand::get?!");
+  llvm_unreachable("Missing Case! Out of sync with OperandOwnership");
 }
 
 SILValue ForwardingOperand::getSingleForwardedValue() const {

lib/SIL/Verifier/SILOwnershipVerifier.cpp

@@ -381,7 +381,7 @@ bool SILValueOwnershipChecker::gatherUsers(
                          << "Address User: " << *op->getUser();
           });
         };
-        foundError |= interiorPointerOperand.getImplicitUses(
+        foundError |= interiorPointerOperand.findTransitiveUses(
             nonLifetimeEndingUsers, &onError);
       }
 

lib/SILOptimizer/SemanticARC/CopyValueOpts.cpp

@@ -331,7 +331,7 @@ static bool canJoinIfCopyDiesInFunctionExitingBlock(
 }
 
 static Operand *lookThroughSingleForwardingUse(Operand *use) {
-  auto forwardingOperand = ForwardingOperand::get(use);
+  ForwardingOperand forwardingOperand(use);
   if (!forwardingOperand)
     return nullptr;
   auto forwardedValue = forwardingOperand.getSingleForwardedValue();
@@ -423,7 +423,7 @@ static bool tryJoinIfDestroyConsumingUseInSameBlock(
     // If not, see if this use did have a forwardedValue but that forwardedValue
     // has multiple end lifetime uses. In that case, we can optimize if there
     // aren't any uses/etc
-    auto forwardingOperand = ForwardingOperand::get(currentForwardingUse);
+    ForwardingOperand forwardingOperand(currentForwardingUse);
     if (!forwardingOperand)
       return false;
     auto forwardedValue = forwardingOperand.getSingleForwardedValue();

lib/SILOptimizer/SemanticARC/OwnershipLiveRange.cpp

@@ -228,7 +228,7 @@ void OwnershipLiveRange::convertOwnedGeneralForwardingUsesToGuaranteed() && {
   while (!ownershipForwardingUses.empty()) {
     auto *use = ownershipForwardingUses.back();
     ownershipForwardingUses = ownershipForwardingUses.drop_back();
-    auto operand = ForwardingOperand::get(use);
+    ForwardingOperand operand(use);
     operand.replaceOwnershipKind(OwnershipKind::Owned,
                                  OwnershipKind::Guaranteed);
   }

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -99,44 +99,71 @@ insertOwnedBaseValueAlongBranchEdge(BranchInst *bi, SILValue innerCopy,
   return phiArg;
 }
 
-static void getAllNonTrivialUsePointsOfBorrowedValue(
+static bool findTransitiveBorrowedUses(
   SILValue value, SmallVectorImpl<Operand *> &usePoints,
   SmallVectorImpl<BorrowingOperand> &reborrowPoints) {
   assert(value.getOwnershipKind() == OwnershipKind::Guaranteed);
 
   unsigned firstOffset = usePoints.size();
-  llvm::copy(value->getUses(), std::back_inserter(usePoints));
-
-  if (usePoints.size() == firstOffset)
-    return;
+  for (Operand *use : value->getUses()) {
+    if (use->getOperandOwnership() != OperandOwnership::NonUse)
+      usePoints.push_back(use);
+  }
 
   // NOTE: Use points resizes in this loop so usePoints.size() may be
   // different every time.
   for (unsigned i = firstOffset; i < usePoints.size(); ++i) {
-    if (auto fOperand = ForwardingOperand::get(usePoints[i])) {
-      fOperand.visitForwardedValues([&](SILValue transitiveValue) {
+    Operand *use = usePoints[i];
+    switch (use->getOperandOwnership()) {
+    case OperandOwnership::NonUse:
+    case OperandOwnership::TrivialUse:
+    case OperandOwnership::ForwardingConsume:
+    case OperandOwnership::DestroyingConsume:
+      llvm_unreachable("this operand cannot handle a guaranteed use");
+
+    case OperandOwnership::ForwardingUnowned:
+    case OperandOwnership::PointerEscape:
+      return false;
+
+    case OperandOwnership::InstantaneousUse:
+    case OperandOwnership::UnownedInstantaneousUse:
+    case OperandOwnership::BitwiseEscape:
+    case OperandOwnership::EndBorrow:
+    case OperandOwnership::Reborrow:
+      break;
+
+    case OperandOwnership::InteriorPointer:
+      // If our base guaranteed value does not have any consuming uses (consider
+      // function arguments), we need to be sure to include interior pointer
+      // operands since we may not get a use from a end_scope instruction.
+      if (!InteriorPointerOperand(use).findTransitiveUses(usePoints)) {
+        return false;
+      }
+      break;
+
+    case OperandOwnership::ForwardingBorrow:
+      ForwardingOperand(use).visitForwardedValues(
+        [&](SILValue transitiveValue) {
           // Do not include transitive uses with 'none' ownership
           if (transitiveValue.getOwnershipKind() == OwnershipKind::None)
             return true;
-        for (auto *transitiveUse : transitiveValue->getUses())
+          for (auto *transitiveUse : transitiveValue->getUses()) {
+            if (transitiveUse->getOperandOwnership()
+                != OperandOwnership::NonUse) {
               usePoints.push_back(transitiveUse);
+            }
+          }
           return true;
         });
-      continue;
-    }
+      break;
 
-    if (auto borrowingOp = BorrowingOperand::get(usePoints[i])) {
-      // If we have a reborrow, we have no further work to do, our reborrow is
-      // already a use and we will handle the reborrow separately.
-      if (borrowingOp.isReborrow())
-        continue;
-
-      // Otherwise, try to grab additional end scope instructions to find more
-      // liveness info. Stash any reborrow uses so that we can eliminate the
-      // reborrow before we are done processing.
-      borrowingOp.visitLocalEndScopeUses([&](Operand *scopeEndingUse) {
-        if (auto scopeEndingBorrowingOp =
-                BorrowingOperand::get(scopeEndingUse)) {
+    case OperandOwnership::Borrow:
+      // Try to grab additional end scope instructions to find more liveness
+      // info. Stash any reborrow uses so that we can eliminate the reborrow
+      // before we are done processing.
+      BorrowingOperand(use).visitLocalEndScopeUses(
+        [&](Operand *scopeEndingUse) {
+          if (auto scopeEndingBorrowingOp = BorrowingOperand(scopeEndingUse)) {
             if (scopeEndingBorrowingOp.isReborrow()) {
               reborrowPoints.push_back(scopeEndingUse);
               return true;
@@ -145,24 +172,15 @@ static void getAllNonTrivialUsePointsOfBorrowedValue(
           usePoints.push_back(scopeEndingUse);
           return true;
         });
-
-      // Now break up all of the reborrows
-      continue;
-    }
-
-    // If our base guaranteed value does not have any consuming uses (consider
-    // function arguments), we need to be sure to include interior pointer
-    // operands since we may not get a use from a end_scope instruction.
-    if (auto intPtrOperand = InteriorPointerOperand::get(usePoints[i])) {
-      intPtrOperand.getImplicitUses(usePoints);
-      continue;
     }
   }
+  return true;
 }
 
-// All callers of our RAUW routines must ensure that their values return true
-// from this.
-static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue) {
+// Determine whether it is valid to replace \p oldValue with \p newValue by
+// directly checking ownership requirements. This does not determine whether the
+// scope of the newValue can be fully extended.
+static bool hasValidRAUWOwnership(SILValue oldValue, SILValue newValue) {
   auto newOwnershipKind = newValue.getOwnershipKind();
 
   // If our new kind is ValueOwnershipKind::None, then we are fine. We
@@ -171,6 +189,18 @@ static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue) {
   if (newOwnershipKind == OwnershipKind::None)
     return true;
 
+  // If our old ownership kind is ValueOwnershipKind::None and our new kind is
+  // not, we may need to do more work that has not been implemented yet. So
+  // bail.
+  //
+  // Due to our requirement that types line up, this can only occur given a
+  // non-trivial typed value with None ownership. This can only happen when
+  // oldValue is a trivial payloaded or no-payload non-trivially typed
+  // enum. That doesn't occur that often so we just bail on it today until we
+  // implement this functionality.
+  if (oldValue.getOwnershipKind() == OwnershipKind::None)
+    return false;
+
   // First check if oldValue is SILUndef. If it is, then we know that:
   //
   // 1. SILUndef (and thus oldValue) must have OwnershipKind::None.
@@ -193,23 +223,26 @@ static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue) {
   if (m->getStage() == SILStage::Raw)
     return false;
 
-  // If our old ownership kind is ValueOwnershipKind::None and our new kind is
-  // not, we may need to do more work that has not been implemented yet. So
-  // bail.
-  //
-  // Due to our requirement that types line up, this can only occur given a
-  // non-trivial typed value with None ownership. This can only happen when
-  // oldValue is a trivial payloaded or no-payload non-trivially typed
-  // enum. That doesn't occur that often so we just bail on it today until we
-  // implement this functionality.
-  auto oldOwnershipKind = SILValue(oldValue).getOwnershipKind();
-  if (oldOwnershipKind != OwnershipKind::None)
   return true;
+}
 
-  // Ok, we have an old ownership kind that is OwnershipKind::None and a new
-  // ownership kind that is not OwnershipKind::None. In that case, for now, do
-  // not perform this transform.
+// Determine whether it is valid to replace \p oldValue with \p newValue and
+// extend the lifetime of \p oldValue to cover the new uses.
+//
+// This updates the OwnershipFixupContext, populating transitiveBorrowedUses and
+// recursiveReborrows.
+static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue,
+                                     OwnershipFixupContext &context) {
+  if (!hasValidRAUWOwnership(oldValue, newValue))
     return false;
+
+  if (oldValue.getOwnershipKind() == OwnershipKind::Guaranteed) {
+    // Check that the old lifetime can be extended and record the necessary
+    // book-keeping in the OwnershipFixupContext.
+    return findTransitiveBorrowedUses(oldValue, context.transitiveBorrowedUses,
+                                      context.recursiveReborrows);
+  }
+  return true;
 }
 
 //===----------------------------------------------------------------------===//
@@ -651,19 +684,17 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleGuaranteed() {
   // Otherwise, we need to actually modify the IR. We first always first
   // lifetime extend newValue to oldValue's transitive uses to set our
   // workspace.
-  SmallVector<Operand *, 32> usePoints;
-  SmallVector<BorrowingOperand, 8> recursiveBorrowScopeReborrows;
-  getAllNonTrivialUsePointsOfBorrowedValue(oldValue, usePoints,
-                                           recursiveBorrowScopeReborrows);
 
   // If we have any transitive reborrows on sub-borrows.
-  if (recursiveBorrowScopeReborrows.size())
-    eliminateReborrowsOfRecursiveBorrows(recursiveBorrowScopeReborrows,
-                                         usePoints, ctx.callbacks);
+  if (ctx.recursiveReborrows.size())
+    eliminateReborrowsOfRecursiveBorrows(ctx.recursiveReborrows,
+                                         ctx.transitiveBorrowedUses,
+                                         ctx.callbacks);
 
   auto extender = getLifetimeExtender();
   SILValue newBorrowedValue =
-      extender.createPlusZeroBorrow<ArrayRef<Operand *>>(newValue, usePoints);
+      extender.createPlusZeroBorrow<ArrayRef<Operand *>>(
+          newValue, ctx.transitiveBorrowedUses);
 
   // Now we need to handle reborrows by eliminating the reborrows from any
   // borrowing operands that use old value as well as from oldvalue itself. We
@@ -696,8 +727,7 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleGuaranteed() {
 
 SILBasicBlock::iterator OwnershipRAUWUtility::perform() {
   assert(oldValue->getFunction()->hasOwnership());
-  assert(
-      canFixUpOwnershipForRAUW(oldValue, newValue) &&
+  assert(hasValidRAUWOwnership(oldValue, newValue) &&
       "Should have checked if can perform this operation before calling it?!");
   // If our new value is just none, we can pass anything to do it so just RAUW
   // and return.
@@ -888,7 +918,7 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
 
   // Otherwise, lets check if we can perform this RAUW operation. If we can't,
   // set ctx to nullptr to invalidate the helper and return.
-  if (!canFixUpOwnershipForRAUW(oldValue, newValue)) {
+  if (!canFixUpOwnershipForRAUW(oldValue, newValue, inputCtx)) {
     ctx = nullptr;
     return;
   }
@@ -964,7 +994,7 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
 
   // For now, just gather up uses
   auto &oldValueUses = ctx->extraAddressFixupInfo.allAddressUsesFromOldValue;
-  if (InteriorPointerOperand::getImplicitUsesForAddress(oldValue,
+  if (InteriorPointerOperand::findTransitiveUsesForAddress(oldValue,
                                                            oldValueUses)) {
     // If we found an error, invalidate and return!
     ctx = nullptr;