Fix findTransitiveUses to only record leaf uses.

Recording uses is now optional. This utility is also used simply to
check for PointerEscape. When uses are recorded, they are used to
extend a live range. There could be a large number of transitive uses
that we don't want to pass back to the caller.

Author: atrick

include/swift/SIL/OwnershipUtils.h

@@ -77,8 +77,9 @@ inline bool isForwardingConsume(SILValue value) {
   return canOpcodeForwardOwnedValues(value);
 }
 
-/// Find all "use points" of \p guaranteedValue that determine its lifetime
-/// requirement.
+/// Find leaf "use points" of \p guaranteedValue that determine its lifetime
+/// requirement. If \p usePoints is nullptr, then the simply returns true if no
+/// PointerEscape use was found.
 ///
 /// Precondition: \p guaranteedValue is not a BorrowedValue.
 ///
@@ -99,10 +100,10 @@ inline bool isForwardingConsume(SILValue value) {
 /// When this is called on a value that does not introduce a new scope, none of
 /// the use points can be EndBorrows or Reborrows. Those uses are only allowed
 /// on borrow-introducing values.
-bool findInnerTransitiveGuaranteedUses(SILValue guaranteedValue,
-                                       SmallVectorImpl<Operand *> &usePoints);
+bool findInnerTransitiveGuaranteedUses(
+    SILValue guaranteedValue, SmallVectorImpl<Operand *> *usePoints = nullptr);
 
-/// Find all "use points" of a guaranteed value within its enclosing borrow
+/// Find leaf "use points" of a guaranteed value within its enclosing borrow
 /// scope (without looking through reborrows). To find the use points of the
 /// extended borrow scope, after looking through reborrows, use
 /// findExtendedTransitiveGuaranteedUses() instead.
@@ -666,7 +667,7 @@ BorrowedValue getSingleBorrowIntroducingValue(SILValue inputValue);
 /// be transitive uses of the given address. Used to implement \see
 /// findTransitiveUses.
 bool findTransitiveUsesForAddress(
-    SILValue address, SmallVectorImpl<Operand *> &foundUses,
+    SILValue address, SmallVectorImpl<Operand *> *foundUses = nullptr,
     std::function<void(Operand *)> *onError = nullptr);
 
 class InteriorPointerOperandKind {
@@ -834,14 +835,17 @@ struct InteriorPointerOperand {
     llvm_unreachable("Covered switch isn't covered?!");
   }
 
-  /// Transitively compute the list of uses that this interior pointer operand
-  /// puts on its parent guaranted value.
+  /// Transitively compute the list of leaf uses that this interior pointer
+  /// operand puts on its parent guaranted value.
+  ///
+  /// If \p foundUses is nullptr, this simply returns true if no PointerEscapes
+  /// were found.
   ///
   /// 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 findTransitiveUses(SmallVectorImpl<Operand *> &foundUses,
+  bool findTransitiveUses(SmallVectorImpl<Operand *> *foundUses = nullptr,
                           std::function<void(Operand *)> *onError = nullptr) {
     return findTransitiveUsesForAddress(getProjectedAddress(), foundUses,
                                         onError);
@@ -913,7 +917,7 @@ struct AddressOwnership {
 
   /// Transitively compute uses of this base address.
   bool findTransitiveUses(SmallVectorImpl<Operand *> &foundUses) {
-    return findTransitiveUsesForAddress(base.getBaseAddress(), foundUses);
+    return findTransitiveUsesForAddress(base.getBaseAddress(), &foundUses);
   }
 
   /// Return true of all \p uses occur before the end of the address' lifetime

lib/SIL/Utils/OwnershipUtils.cpp

@@ -12,6 +12,7 @@
 
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/Basic/Defer.h"
+#include "swift/Basic/DAGNodeWorklist.h"
 #include "swift/Basic/SmallPtrSetVector.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/LinearLifetimeChecker.h"
@@ -160,35 +161,33 @@ bool swift::canOpcodeForwardOwnedValues(Operand *use) {
 // points. Transitively find all nested scope-ending instructions by looking
 // through nested reborrows. Nested reborrows are not use points and \p
 // visitReborrow is not called for them.
-bool swift::
-findInnerTransitiveGuaranteedUses(SILValue guaranteedValue,
-                                  SmallVectorImpl<Operand *> &usePoints) {
+bool swift::findInnerTransitiveGuaranteedUses(
+    SILValue guaranteedValue, SmallVectorImpl<Operand *> *usePoints) {
   // Push the value's immediate uses.
-  unsigned firstOffset = usePoints.size();
+  //
+  // TODO: The worklist can be a simple vector without any a membership check if
+  // destructures are changed to be represented as reborrows. Currently a
+  // destructure forwards multiple results! This means that usePoints could grow
+  // exponentially without the membership check. It's fine to do this membership
+  // check locally in this function (within a borrow scope) because it isn't
+  // needed for the immediate uses, only the transitive uses.
+  DAGNodeWorklist<Operand *, 8> worklist;
   for (Operand *use : guaranteedValue->getUses()) {
     if (use->getOperandOwnership() != OperandOwnership::NonUse)
-      usePoints.push_back(use);
+      worklist.insert(use);
   }
 
   // --- Transitively follow forwarded uses and look for escapes.
 
-  // TODO: Remove this SmallPtrSet if destructures are changed to be represented
-  // as reborrows. Currently it forwards multiple results! This means that
-  // usePoints could grow exponentially without a membership check. It's fine to
-  // do this membership check locally in this function (within a borrow
-  // scope). It isn't needed for the immediate uses, only the transitive uses.
-  SmallPtrSet<Operand *, 16> visitedUses;
-  auto pushUse = [&](Operand *use) {
-    if (use->getOperandOwnership() != OperandOwnership::NonUse) {
-      if (visitedUses.insert(use).second)
-        usePoints.push_back(use);
+  auto recordLeafUse = [&](Operand *use) {
+    if (usePoints && use->getOperandOwnership() != OperandOwnership::NonUse) {
+      usePoints->push_back(use);
     }
     return true;
   };
 
   // usePoints grows in this loop.
-  for (unsigned i = firstOffset; i < usePoints.size(); ++i) {
-    Operand *use = usePoints[i];
+  while (Operand *use = worklist.pop()) {
     switch (use->getOperandOwnership()) {
     case OperandOwnership::NonUse:
     case OperandOwnership::TrivialUse:
@@ -228,14 +227,17 @@ findInnerTransitiveGuaranteedUses(SILValue guaranteedValue,
             if (transitiveValue.getOwnershipKind() == OwnershipKind::None)
               return true;
             for (auto *transitiveUse : transitiveValue->getUses()) {
-              pushUse(transitiveUse);
+              if (transitiveUse->getOperandOwnership()
+                  != OperandOwnership::NonUse) {
+                worklist.insert(use);
+              }
             }
             return true;
           });
       break;
 
     case OperandOwnership::Borrow:
-      BorrowingOperand(use).visitExtendedScopeEndingUses(pushUse);
+      BorrowingOperand(use).visitExtendedScopeEndingUses(recordLeafUse);
     }
   }
   return true;
@@ -271,7 +273,7 @@ bool swift::findTransitiveGuaranteedUses(
     }
     return true;
   }
-  return findInnerTransitiveGuaranteedUses(guaranteedValue, usePoints);
+  return findInnerTransitiveGuaranteedUses(guaranteedValue, &usePoints);
 }
 
 // Find all use points of \p guaranteedValue within its borrow scope. If the
@@ -705,23 +707,24 @@ bool BorrowedValue::visitInteriorPointerOperandHelper(
 }
 
 bool swift::findTransitiveUsesForAddress(
-    SILValue projectedAddress, SmallVectorImpl<Operand *> &foundUses,
+    SILValue projectedAddress, SmallVectorImpl<Operand *> *foundUses,
     std::function<void(Operand *)> *onError) {
   SmallVector<Operand *, 8> worklist(projectedAddress->getUses());
 
   bool foundError = false;
 
+  auto leafUse = [foundUses](Operand *use) {
+    if (foundUses)
+      foundUses->push_back(use);
+  };
+
   while (!worklist.empty()) {
     auto *op = worklist.pop_back_val();
 
     // Skip type dependent operands.
     if (op->isTypeDependent())
       continue;
 
-    // Before we do anything, add this operand to our implicit regular user
-    // list.
-    foundUses.push_back(op);
-
     // Then update the worklist with new things to find if we recognize this
     // inst and then continue. If we fail, we emit an error at the bottom of the
     // loop that we didn't recognize the user.
@@ -741,6 +744,7 @@ bool swift::findTransitiveUsesForAddress(
         isa<SwitchEnumAddrInst>(user) || isa<CheckedCastAddrBranchInst>(user) ||
         isa<SelectEnumAddrInst>(user) || isa<InjectEnumAddrInst>(user) ||
         isa<IsUniqueInst>(user)) {
+      leafUse(op);
       continue;
     }
 
@@ -763,33 +767,39 @@ bool swift::findTransitiveUsesForAddress(
     if (auto *builtin = dyn_cast<BuiltinInst>(user)) {
       if (auto kind = builtin->getBuiltinKind()) {
         if (*kind == BuiltinValueKind::TSanInoutAccess) {
+          leafUse(op);
           continue;
         }
       }
     }
 
     // If we have a load_borrow, add it's end scope to the liveness requirement.
     if (auto *lbi = dyn_cast<LoadBorrowInst>(user)) {
-      transform(lbi->getEndBorrows(), std::back_inserter(foundUses),
-                [](EndBorrowInst *ebi) { return &ebi->getAllOperands()[0]; });
+      if (foundUses) {
+        transform(lbi->getEndBorrows(), std::back_inserter(*foundUses),
+                  [](EndBorrowInst *ebi) { return &ebi->getAllOperands()[0]; });
+      }
       continue;
     }
 
     // TODO: Merge this into the full apply site code below.
     if (auto *beginApply = dyn_cast<BeginApplyInst>(user)) {
-      // TODO: Just add this to implicit regular user list?
-      llvm::copy(beginApply->getTokenResult()->getUses(),
-                 std::back_inserter(foundUses));
+      if (foundUses) {
+        llvm::copy(beginApply->getTokenResult()->getUses(),
+                   std::back_inserter(*foundUses));
+      }
       continue;
     }
 
     if (auto fas = FullApplySite::isa(user)) {
+      leafUse(op);
       continue;
     }
 
     if (auto *mdi = dyn_cast<MarkDependenceInst>(user)) {
       // If this is the base, just treat it as a liveness use.
       if (op->get() == mdi->getBase()) {
+        leafUse(op);
         continue;
       }
 

lib/SIL/Verifier/SILOwnershipVerifier.cpp

@@ -377,7 +377,7 @@ bool SILValueOwnershipChecker::gatherUsers(
           });
         };
         foundError |= !interiorPointerOperand.findTransitiveUses(
-            nonLifetimeEndingUsers, &onError);
+            &nonLifetimeEndingUsers, &onError);
       }
 
       // Finally add the op to the non lifetime ending user list.

lib/SILOptimizer/Utils/BasicBlockOptUtils.cpp

@@ -100,8 +100,7 @@ static bool canBorrowGuaranteedResult(SILValue guaranteedResult) {
     // conversion to a non-guaranteed value. Either way, not interesting.
     return true;
   }
-  SmallVector<Operand *, 16> usePoints;
-  return findInnerTransitiveGuaranteedUses(guaranteedResult, usePoints);
+  return findInnerTransitiveGuaranteedUses(guaranteedResult);
 }
 
 bool swift::canCloneTerminator(TermInst *termInst) {

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -1150,7 +1150,7 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
   auto &oldValueUses = ctx->extraAddressFixupInfo.allAddressUsesFromOldValue;
   // FIXME: The return value of findTransitiveUsesForAddress is currently
   // inverted.
-  if (findTransitiveUsesForAddress(oldValue, oldValueUses)) {
+  if (!findTransitiveUsesForAddress(oldValue, &oldValueUses)) {
     invalidate();
     return;
   }
@@ -1520,7 +1520,7 @@ void GuaranteedPhiBorrowFixup::insertEndBorrowsAndFindPhis(
     return;
   }
   SmallVector<Operand *, 16> usePoints;
-  bool result = findInnerTransitiveGuaranteedUses(phi, usePoints);
+  bool result = findInnerTransitiveGuaranteedUses(phi, &usePoints);
   assert(result && "should be checked by canCloneTerminator");
   (void)result;