CanonicalOSSALifetime: rewrite consolidateBorrowScope

It now handles looking through forwarding consumes such as
destructures.

Expected to be NFC since borrow consolidation is still disabled by
default.

TODO: Write unit tests for various forwarding consumes in addition to
destructure.

Author: atrick

include/swift/SILOptimizer/Utils/CanonicalOSSALifetime.h

@@ -213,6 +213,50 @@ class CanonicalOSSAConsumeInfo {
   SWIFT_ASSERT_ONLY_DECL(void dump() const LLVM_ATTRIBUTE_USED);
 };
 
+// Worklist of pointer-like things that have an invalid default value. Avoid
+// revisiting nodes--suitable for DAGs, but pops finished nodes without
+// preserving them in the vector.
+//
+// The primary API has two methods: intialize() and pop(). Others are provided
+// for flexibility.
+//
+// TODO: make this a better utility.
+template <typename T, unsigned SmallSize> struct PtrWorklist {
+  SmallPtrSet<T, SmallSize> ptrVisited;
+  SmallVector<T, SmallSize> ptrVector;
+
+  PtrWorklist() = default;
+
+  PtrWorklist(const PtrWorklist &) = delete;
+
+  void initialize(T t) {
+    clear();
+    insert(t);
+  }
+
+  template <typename R> void initializeRange(R &&range) {
+    clear();
+    ptrVisited.insert(range.begin(), range.end());
+    ptrVector.append(range.begin(), range.end());
+  }
+
+  T pop() { return empty() ? T() : ptrVector.pop_back_val(); }
+
+  bool empty() const { return ptrVector.empty(); }
+
+  unsigned size() const { return ptrVector.size(); }
+
+  void clear() {
+    ptrVector.clear();
+    ptrVisited.clear();
+  }
+
+  void insert(T t) {
+    if (ptrVisited.insert(t).second)
+      ptrVector.push_back(t);
+  }
+};
+
 /// Canonicalize OSSA lifetimes.
 ///
 /// Allows the allocation of analysis state to be reused across calls to
@@ -227,6 +271,10 @@ class CanonicalizeOSSALifetime {
   /// liveness may be pruned during canonicalization.
   bool pruneDebugMode;
 
+  /// If true, borrows scopes will be canonicalized allowing copies of
+  /// guaranteed values to be eliminated.
+  bool canonicalizeBorrowMode;
+
   /// If true, then new destroy_value instructions will be poison.
   bool poisonRefsMode;
 
@@ -237,8 +285,6 @@ class CanonicalizeOSSALifetime {
 
   DominanceAnalysis *dominanceAnalysis;
 
-  DeadEndBlocks *deBlocks;
-
   /// Current copied def for which this state describes the liveness.
   SILValue currentDef;
 
@@ -262,11 +308,11 @@ class CanonicalizeOSSALifetime {
   /// outisde the pruned liveness at the time it is discovered.
   llvm::SmallPtrSet<DebugValueInst *, 8> debugValues;
 
-  /// Reuse a general worklist for def-use traversal.
-  SmallSetVector<SILValue, 8> defUseWorklist;
+  /// Reuse a general visited set for def-use traversal.
+  PtrWorklist<SILValue, 8> defUseWorklist;
 
   /// Reuse a general worklist for CFG traversal.
-  SmallSetVector<SILBasicBlock *, 8> blockWorklist;
+  PtrWorklist<SILBasicBlock *, 8> blockWorklist;
 
   /// Pruned liveness for the extended live range including copies. For this
   /// purpose, only consuming instructions are considered "lifetime
@@ -288,13 +334,15 @@ class CanonicalizeOSSALifetime {
   CanonicalOSSAConsumeInfo consumes;
 
 public:
-  CanonicalizeOSSALifetime(bool pruneDebugMode, bool poisonRefsMode,
+  CanonicalizeOSSALifetime(bool pruneDebugMode, bool canonicalizeBorrowMode,
+                           bool poisonRefsMode,
                            NonLocalAccessBlockAnalysis *accessBlockAnalysis,
-                           DominanceAnalysis *dominanceAnalysis,
-                           DeadEndBlocks *deBlocks)
-      : pruneDebugMode(pruneDebugMode), poisonRefsMode(poisonRefsMode),
-        accessBlockAnalysis(accessBlockAnalysis),
-        dominanceAnalysis(dominanceAnalysis), deBlocks(deBlocks) {}
+                           DominanceAnalysis *dominanceAnalysis)
+    : pruneDebugMode(pruneDebugMode),
+      canonicalizeBorrowMode(canonicalizeBorrowMode),
+      poisonRefsMode(poisonRefsMode),
+      accessBlockAnalysis(accessBlockAnalysis),
+      dominanceAnalysis(dominanceAnalysis) {}
 
   SILValue getCurrentDef() const { return currentDef; }
 
@@ -351,6 +399,15 @@ class CanonicalizeOSSALifetime {
 
   bool consolidateBorrowScope();
 
+  bool findBorrowScopeUses(llvm::SmallPtrSetImpl<SILInstruction *> &useInsts);
+
+  void filterOuterBorrowUseInsts(
+      llvm::SmallPtrSetImpl<SILInstruction *> &outerUseInsts);
+
+  void rewriteOuterBorrowUsesAndFindConsumes(
+      SILValue incomingValue,
+      llvm::SmallPtrSetImpl<SILInstruction *> &outerUseInsts);
+
   bool computeCanonicalLiveness();
 
   bool endsAccessOverlappingPrunedBoundary(SILInstruction *inst);

lib/SILOptimizer/Transforms/CopyPropagation.cpp

@@ -26,6 +26,7 @@
 #define DEBUG_TYPE "copy-propagation"
 
 #include "swift/SIL/BasicBlockUtils.h"
+#include "swift/SIL/SILUndef.h"
 #include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
@@ -34,6 +35,12 @@
 
 using namespace swift;
 
+// This only applies to -O copy-propagation.
+llvm::cl::opt<bool>
+    EnableRewriteBorrows("canonical-ossa-rewrite-borrows",
+                         llvm::cl::init(false),
+                         llvm::cl::desc("Enable rewriting borrow scopes"));
+
 //===----------------------------------------------------------------------===//
 // CopyPropagation: Top-Level Function Transform.
 //===----------------------------------------------------------------------===//
@@ -44,13 +51,17 @@ class CopyPropagation : public SILFunctionTransform {
   bool pruneDebug;
   /// True of all values should be canonicalized.
   bool canonicalizeAll;
+  /// If true, then borrow scopes will be canonicalized, allowing copies of
+  /// guaranteed values to be optimized. Does *not* shrink the borrow scope.
+  bool canonicalizeBorrows;
   /// If true, then new destroy_value instructions will be poison.
   bool poisonRefs;
 
 public:
-  CopyPropagation(bool pruneDebug, bool canonicalizeAll, bool poisonRefs)
-    : pruneDebug(pruneDebug), canonicalizeAll(canonicalizeAll),
-      poisonRefs(poisonRefs) {}
+  CopyPropagation(bool pruneDebug, bool canonicalizeAll,
+                  bool canonicalizeBorrows, bool poisonRefs)
+      : pruneDebug(pruneDebug), canonicalizeAll(canonicalizeAll),
+        canonicalizeBorrows(canonicalizeBorrows), poisonRefs(poisonRefs) {}
 
   /// The entry point to this function transformation.
   void run() override;
@@ -71,12 +82,18 @@ static bool isCopyDead(CopyValueInst *copy, bool pruneDebug) {
   return true;
 }
 
+static bool isBorrowDead(BeginBorrowInst *borrow) {
+  return llvm::all_of(borrow->getUses(), [](Operand *use) {
+      SILInstruction *user = use->getUser();
+      return isa<EndBorrowInst>(user) || user->isDebugInstruction();
+    });
+}
+
 /// Top-level pass driver.
 void CopyPropagation::run() {
   auto *f = getFunction();
   auto *accessBlockAnalysis = getAnalysis<NonLocalAccessBlockAnalysis>();
   auto *dominanceAnalysis = getAnalysis<DominanceAnalysis>();
-  auto *deBlocksAnalysis = getAnalysis<DeadEndBlocksAnalysis>();
 
   // Debug label for unit testing.
   LLVM_DEBUG(llvm::dbgs() << "*** CopyPropagation: " << f->getName() << "\n");
@@ -115,6 +132,12 @@ void CopyPropagation::run() {
         || SILValue(extract).getOwnershipKind() != OwnershipKind::Guaranteed)
       continue;
 
+    // Bail-out if we won't rewrite borrows because that currently regresses
+    // Breadcrumbs.MutatedUTF16ToIdx.Mixed/Breadcrumbs.MutatedIdxToUTF16.Mixed.
+    // Also, mandatory copyprop does not need to rewrite destructures.
+    if (!canonicalizeBorrows)
+      continue;
+
     if (SILValue destructuredResult = convertExtractToDestructure(extract)) {
       // Remove to-be-deleted instructions from copiedDeds. The extract cannot
       // be in the copiedDefs set since getCanonicalCopiedDef does not allow a
@@ -138,14 +161,13 @@ void CopyPropagation::run() {
     }
   }
   // Perform copy propgation for each copied value.
-  CanonicalizeOSSALifetime canonicalizer(pruneDebug, poisonRefs,
-                                         accessBlockAnalysis,
-                                         dominanceAnalysis,
-                                         deBlocksAnalysis->get(f));
+  CanonicalizeOSSALifetime canonicalizer(pruneDebug, canonicalizeBorrows,
+                                         poisonRefs, accessBlockAnalysis,
+                                         dominanceAnalysis);
   // Cleanup dead copies. If getCanonicalCopiedDef returns a copy (because the
-  // copy's source operand is unrecgonized), then thecan copy is itself treated
+  // copy's source operand is unrecgonized), then the copy is itself treated
   // like a def and may be dead after canonicalization.
-  llvm::SmallVector<CopyValueInst *, 4> deadCopies;
+  llvm::SmallVector<SILInstruction *, 4> deadCopies;
   for (auto &def : copiedDefs) {
     // Canonicalized this def.
     canonicalizer.canonicalizeValueLifetime(def);
@@ -155,6 +177,14 @@ void CopyPropagation::run() {
         deadCopies.push_back(copy);
       }
     }
+    // Dead borrow scopes must be removed as uses before canonicalizing the
+    // outer copy.
+    if (auto *borrow = dyn_cast<BeginBorrowInst>(def)) {
+      if (isBorrowDead(borrow)) {
+        borrow->setOperand(SILUndef::get(borrow->getType(), *f));
+        deadCopies.push_back(borrow);
+      }
+    }
     // Canonicalize any new outer copy.
     if (SILValue outerCopy = canonicalizer.createdOuterCopy()) {
       SILValue outerDef = canonicalizer.getCanonicalCopiedDef(outerCopy);
@@ -165,26 +195,29 @@ void CopyPropagation::run() {
   }
   if (canonicalizer.hasChanged() || !deadCopies.empty()) {
     InstructionDeleter deleter;
-    for (auto *copy : deadCopies) {
-      deleter.recursivelyDeleteUsersIfDead(copy);
+    for (auto *copyOrBorrow : deadCopies) {
+      deleter.recursivelyDeleteUsersIfDead(copyOrBorrow);
     }
     // Preserves NonLocalAccessBlockAnalysis.
     accessBlockAnalysis->lockInvalidation();
     invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
     accessBlockAnalysis->unlockInvalidation();
     if (f->getModule().getOptions().VerifySILOwnership) {
+      auto *deBlocksAnalysis = getAnalysis<DeadEndBlocksAnalysis>();
       f->verifyOwnership(deBlocksAnalysis->get(f));
     }
   }
 }
 
 SILTransform *swift::createMandatoryCopyPropagation() {
   return new CopyPropagation(/*pruneDebug*/ true, /*canonicalizeAll*/ true,
+                             /*canonicalizeBorrows*/ false,
                              /*poisonRefs*/ true);
 }
 
 SILTransform *swift::createCopyPropagation() {
   return new CopyPropagation(/*pruneDebug*/ true, /*canonicalizeAll*/ false,
+                             /*canonicalizeBorrows*/ EnableRewriteBorrows,
                              /*poisonRefs*/ false);
 }