Merge pull request swiftlang#40598 from nate-chandler/lexical_lifetimes/shrink-borrow-scope/hoist-over-copy-users

[CopyPropagation] Hoist end_borrows over users of copies.

Author: nate-chandler

include/swift/SIL/OwnershipUtils.h

@@ -104,12 +104,9 @@ inline bool isForwardingConsume(SILValue value) {
 bool findInnerTransitiveGuaranteedUses(
     SILValue guaranteedValue, SmallVectorImpl<Operand *> *usePoints = nullptr);
 
-/// Like findInnerTransitiveGuaranteedUses except that rather than it being a
-/// precondition that the provided value not be a BorrowedValue, it is a [type-
-/// system-enforced] precondition that the provided value be a BorrowedValue.
-///
-/// TODO: Merge with findInnerTransitiveGuaranteedUses.
-bool findInnerTransitiveGuaranteedUsesOfBorrowedValue(
+/// Find all uses in the extended lifetime (i.e. including copies) of a simple
+/// (i.e. not reborrowed) borrow scope and its transitive uses.
+bool findExtendedUsesOfSimpleBorrowedValue(
     BorrowedValue borrowedValue,
     SmallVectorImpl<Operand *> *usePoints = nullptr);
 

include/swift/SIL/SILValue.h

@@ -438,7 +438,7 @@ class ValueBase : public SILNode, public SILAllocated<ValueBase> {
 
   template <class T> inline T *getSingleConsumingUserOfType() const;
 
-  /// Returns true if this operand has exactly two.
+  /// Returns true if this operand has exactly two uses.
   ///
   /// This is useful if one has found a predefined set of 2 unique users and
   /// wants to check if there are any other users without iterating over the

include/swift/SILOptimizer/Utils/CanonicalizeBorrowScope.h

@@ -147,7 +147,9 @@ class CanonicalizeBorrowScope {
   bool consolidateBorrowScope();
 };
 
-bool shrinkBorrowScope(BeginBorrowInst *bbi, InstructionDeleter &deleter);
+bool shrinkBorrowScope(
+    BeginBorrowInst *bbi, InstructionDeleter &deleter,
+    SmallVectorImpl<CopyValueInst *> &modifiedCopyValueInsts);
 
 } // namespace swift
 

lib/SIL/Utils/OwnershipUtils.cpp

@@ -195,14 +195,9 @@ bool swift::findInnerTransitiveGuaranteedUses(
   return true;
 }
 
-/// Like findInnerTransitiveGuaranteedUses except that rather than it being a
-/// precondition that the provided value not be a BorrowedValue, it is a [type-
-/// system-enforced] precondition that the provided value be a BorrowedValue.
-///
-/// TODO: Merge with findInnerTransitiveGuaranteedUses.  Note that at the moment
-///       the two are _almost_ identical, but not quite because the other has a
-///       #if 0 and not just leaf uses but ALL uses are recorded.
-bool swift::findInnerTransitiveGuaranteedUsesOfBorrowedValue(
+/// Find all uses in the extended lifetime (i.e. including copies) of a simple
+/// (i.e. not reborrowed) borrow scope and its transitive uses.
+bool swift::findExtendedUsesOfSimpleBorrowedValue(
     BorrowedValue borrowedValue, SmallVectorImpl<Operand *> *usePoints) {
 
   auto recordUse = [&](Operand *use) {
@@ -220,21 +215,31 @@ bool swift::findInnerTransitiveGuaranteedUsesOfBorrowedValue(
   // membership check locally in this function (within a borrow scope) because
   // it isn't needed for the immediate uses, only the transitive uses.
   GraphNodeWorklist<Operand *, 8> worklist;
-  for (Operand *use : borrowedValue.value->getUses()) {
-    if (use->getOperandOwnership() != OperandOwnership::NonUse)
-      worklist.insert(use);
-  }
+  auto addUsesToWorklist = [&worklist](SILValue value) {
+    for (Operand *use : value->getUses()) {
+      if (use->getOperandOwnership() != OperandOwnership::NonUse)
+        worklist.insert(use);
+    }
+  };
+
+  addUsesToWorklist(borrowedValue.value);
 
   // --- Transitively follow forwarded uses and look for escapes.
 
   // usePoints grows in this loop.
   while (Operand *use = worklist.pop()) {
+    if (auto *cvi = dyn_cast<CopyValueInst>(use->getUser())) {
+      addUsesToWorklist(cvi);
+    }
     switch (use->getOperandOwnership()) {
     case OperandOwnership::NonUse:
+      break;
+
     case OperandOwnership::TrivialUse:
     case OperandOwnership::ForwardingConsume:
     case OperandOwnership::DestroyingConsume:
-      llvm_unreachable("this operand cannot handle an inner guaranteed use");
+      recordUse(use);
+      break;
 
     case OperandOwnership::ForwardingUnowned:
     case OperandOwnership::PointerEscape:
@@ -264,6 +269,7 @@ bool swift::findInnerTransitiveGuaranteedUsesOfBorrowedValue(
           AddressUseKind::NonEscaping) {
         return false;
       }
+      recordUse(use);
       break;
 
     case OperandOwnership::ForwardingBorrow: {

lib/SILOptimizer/Transforms/CopyPropagation.cpp

@@ -438,8 +438,12 @@ void CopyPropagation::run() {
   //       blocks and pushing begin_borrows as we see them and then popping them
   //       off the end will result in shrinking inner borrow scopes first.
   while (auto *bbi = beginBorrowsToShrink.pop()) {
-    changed |= shrinkBorrowScope(bbi, deleter);
+    SmallVector<CopyValueInst *, 4> modifiedCopyValueInsts;
+    changed |= shrinkBorrowScope(bbi, deleter, modifiedCopyValueInsts);
+    for (auto *cvi : modifiedCopyValueInsts)
+      defWorklist.updateForCopy(cvi);
   }
+  deleter.cleanupDeadInstructions();
 
   // canonicalizer performs all modifications through deleter's callbacks, so we
   // don't need to explicitly check for changes.