Fix -canonical-ossa-rewrite-borrows but leave it disabled.

include/swift/SIL/OwnershipUtils.h

@@ -217,6 +217,9 @@ struct BorrowingOperand {
   /// over a region of code instead of just for a single instruction, visit
   /// those uses.
   ///
+  /// Returns true if all visitor invocations returns true. Exits early if a
+  /// visitor returns false.
+  ///
   /// Example: An apply performs an instantaneous recursive borrow of a
   /// guaranteed value but a begin_apply borrows the value over the entire
   /// region of code corresponding to the coroutine.

include/swift/SILOptimizer/Utils/CanonicalOSSALifetime.h

@@ -196,6 +196,8 @@ class CanonicalizeOSSALifetime {
 
   DominanceAnalysis *dominanceAnalysis;
 
+  DeadEndBlocks *deBlocks;
+
   /// Current copied def for which this state describes the liveness.
   SILValue currentDef;
 
@@ -237,9 +239,10 @@ class CanonicalizeOSSALifetime {
 public:
   CanonicalizeOSSALifetime(bool pruneDebug,
                            NonLocalAccessBlockAnalysis *accessBlockAnalysis,
-                           DominanceAnalysis *dominanceAnalysis)
+                           DominanceAnalysis *dominanceAnalysis,
+                           DeadEndBlocks *deBlocks)
     : pruneDebug(pruneDebug), accessBlockAnalysis(accessBlockAnalysis),
-      dominanceAnalysis(dominanceAnalysis) {}
+      dominanceAnalysis(dominanceAnalysis), deBlocks(deBlocks) {}
 
   SILValue getCurrentDef() const { return currentDef; }
 
@@ -293,7 +296,7 @@ class CanonicalizeOSSALifetime {
 
   bool computeBorrowLiveness();
 
-  void consolidateBorrowScope();
+  bool consolidateBorrowScope();
 
   bool computeCanonicalLiveness();
 

lib/SILOptimizer/Transforms/CopyPropagation.cpp

@@ -26,9 +26,11 @@
 #define DEBUG_TYPE "copy-propagation"
 
 #include "swift/SIL/BasicBlockUtils.h"
+#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/CanonicalOSSALifetime.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 
 using namespace swift;
 
@@ -49,11 +51,26 @@ class CopyPropagation : public SILFunctionTransform {
 };
 } // end anonymous namespace
 
+static bool isCopyDead(CopyValueInst *copy, bool pruneDebug) {
+  for (Operand *use : copy->getUses()) {
+    auto *user = use->getUser();
+    if (isa<DestroyValueInst>(user)) {
+      continue;
+    }
+    if (pruneDebug && isa<DebugValueInst>(user)) {
+      continue;
+    }
+    return false;
+  }
+  return true;
+}
+
 /// 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");
@@ -73,23 +90,39 @@ void CopyPropagation::run() {
   }
   // Perform copy propgation for each copied value.
   CanonicalizeOSSALifetime canonicalizer(pruneDebug, accessBlockAnalysis,
-                                         dominanceAnalysis);
+                                         dominanceAnalysis,
+                                         deBlocksAnalysis->get(f));
+  // Cleanup dead copies. If getCanonicalCopiedDef returns a copy (because the
+  // 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;
   for (auto &def : copiedDefs) {
+    // Canonicalized this def.
     canonicalizer.canonicalizeValueLifetime(def);
+
+    if (auto *copy = dyn_cast<CopyValueInst>(def)) {
+      if (isCopyDead(copy, pruneDebug)) {
+        deadCopies.push_back(copy);
+      }
+    }
+    // Canonicalize any new outer copy.
     if (SILValue outerCopy = canonicalizer.createdOuterCopy()) {
       SILValue outerDef = canonicalizer.getCanonicalCopiedDef(outerCopy);
       canonicalizer.canonicalizeValueLifetime(outerDef);
     }
     // TODO: also canonicalize any lifetime.persistentCopies like separate owned
     // live ranges.
   }
-  if (canonicalizer.hasChanged()) {
+  if (canonicalizer.hasChanged() || !deadCopies.empty()) {
+    InstructionDeleter deleter;
+    for (auto *copy : deadCopies) {
+      deleter.recursivelyDeleteUsersIfDead(copy);
+    }
     // Preserves NonLocalAccessBlockAnalysis.
     accessBlockAnalysis->lockInvalidation();
     invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
     accessBlockAnalysis->unlockInvalidation();
-    DeadEndBlocks deBlocks(f);
-    f->verifyOwnership(&deBlocks);
+    f->verifyOwnership(deBlocksAnalysis->get(f));
   }
 }
 

lib/SILOptimizer/Utils/CanonicalOSSALifetime.cpp

@@ -45,6 +45,7 @@
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
+#include "swift/SILOptimizer/Utils/ValueLifetime.h"
 #include "llvm/ADT/Statistic.h"
 
 using namespace swift;
@@ -73,8 +74,19 @@ SILValue CanonicalizeOSSALifetime::getCanonicalCopiedDef(SILValue v) {
       case BorrowedValueKind::Invalid:
         llvm_unreachable("Using invalid case?!");
       case BorrowedValueKind::SILFunctionArgument:
-      case BorrowedValueKind::BeginBorrow:
         return def;
+      case BorrowedValueKind::BeginBorrow: {
+        bool localScope = true;
+        // TODO: visitLocalScopeEndingUses should have an early exit.
+        borrowedVal.visitLocalScopeEndingUses([&](Operand *endBorrow) {
+          if (endBorrow->getUser()->getParent() != def->getParentBlock())
+            localScope = false;
+        });
+        if (localScope) {
+          return def;
+        }
+        break;
+      }
       case BorrowedValueKind::LoadBorrow:
       case BorrowedValueKind::Phi:
         break;
@@ -87,6 +99,20 @@ SILValue CanonicalizeOSSALifetime::getCanonicalCopiedDef(SILValue v) {
   return v;
 }
 
+/// The lifetime extends beyond given consuming use. Copy the value.
+static void copyLiveUse(Operand *use) {
+  SILInstruction *user = use->getUser();
+  SILBuilderWithScope B(user->getIterator());
+
+  auto loc = RegularLocation::getAutoGeneratedLocation(
+    user->getLoc().getSourceLoc());
+  auto *copy = B.createCopyValue(loc, use->get());
+  use->set(copy);
+
+  ++NumCopiesGenerated;
+  LLVM_DEBUG(llvm::dbgs() << "  Copying at last use " << *copy);
+}
+
 //===----------------------------------------------------------------------===//
 //                        MARK: Rewrite borrow scopes
 //===----------------------------------------------------------------------===//
@@ -121,12 +147,11 @@ bool CanonicalizeOSSALifetime::computeBorrowLiveness() {
     liveness.updateForUse(use->getUser(), /*lifetimeEnding*/ true);
   });
 
-  // TODO: Remove this check. Canonicalize multi-block borrow scopes only after
-  // consolidateBorrowScope can handle persistentCopies, otherwise we may end up
-  // generating more dynamic copies than the non-canonical form.
-  if (liveness.numLiveBlocks() > 1) {
-    return false;
-  }
+  // TODO: Fix getCanonicalCopiedDef to allow multi-block borrows and remove
+  // this assert. This should only be done once consolidateBorrowScope can
+  // handle persistentCopies, otherwise we may end up generating more dynamic
+  // copies than the non-canonical form.
+  assert(liveness.numLiveBlocks() == 1);
   return true;
 }
 
@@ -153,15 +178,26 @@ static CopyValueInst *createOuterCopy(BeginBorrowInst *beginBorrow) {
 // TODO: Canonicalize multi-block borrow scopes, load_borrow scope, and phi
 // borrow scopes by adding one copy per block to persistentCopies for
 // each block that dominates an outer use.
-void CanonicalizeOSSALifetime::consolidateBorrowScope() {
+bool CanonicalizeOSSALifetime::consolidateBorrowScope() {
   if (isa<SILFunctionArgument>(currentDef)) {
-    return;
+    return true;
   }
   // Gather all outer uses before rewriting any to avoid scanning any basic
   // block more than once.
   SmallVector<Operand *, 8> outerUses;
   llvm::SmallPtrSet<SILInstruction *, 8> outerUseInsts;
+  auto isUserInLiveOutBlock = [&](SILInstruction *user) {
+    // TODO: enable isUserInLiveOutBlock once we support multi-block borrows
+    // return (liveness.getBlockLiveness(user->getParent())
+    //        == PrunedLiveBlocks::LiveOut);
+    return false;
+  };
   auto recordOuterUse = [&](Operand *use) {
+    // If the user's block is LiveOut, then it is definitely within the borrow
+    // scope, so there's no need to record it.
+    if (isUserInLiveOutBlock(use->getUser())) {
+      return;
+    }
     outerUses.push_back(use);
     outerUseInsts.insert(use->getUser());
   };
@@ -177,12 +213,10 @@ void CanonicalizeOSSALifetime::consolidateBorrowScope() {
         defUseWorklist.insert(copy);
         continue;
       }
-      // debug_value uses are handled like normal uses here. They should be
-      // stripped later if required when handling outerCopy or persistentCopies.
-      if (liveness.getBlockLiveness(user->getParent())
-          == PrunedLiveBlocks::LiveOut) {
-        continue;
-      }
+      // Note: debug_value uses are handled like normal uses here. They should
+      // be stripped later if required when handling outerCopy or
+      // persistentCopies.
+
       switch (use->getOperandOwnership()) {
       case OperandOwnership::NonUse:
         break;
@@ -198,41 +232,102 @@ void CanonicalizeOSSALifetime::consolidateBorrowScope() {
 
       case OperandOwnership::ForwardingUnowned:
       case OperandOwnership::PointerEscape:
+        return false;
+
       case OperandOwnership::InstantaneousUse:
       case OperandOwnership::UnownedInstantaneousUse:
       case OperandOwnership::BitwiseEscape:
       case OperandOwnership::ForwardingConsume:
       case OperandOwnership::DestroyingConsume:
+        recordOuterUse(use);
+        break;
       case OperandOwnership::Borrow:
+        BorrowingOperand borrowOper(use);
+        if (borrowOper.kind == BorrowingOperandKind::Invalid) {
+          return false;
+        }
         recordOuterUse(use);
+        // For borrows, record the scope-ending instructions in addition to the
+        // borrow instruction as an outer use point.
+        borrowOper.visitLocalEndScopeUses([&](Operand *endBorrow) {
+          if (!isUserInLiveOutBlock(endBorrow->getUser())) {
+            outerUseInsts.insert(endBorrow->getUser());
+          }
+          return true;
+        });
         break;
       }
     }
   } // end def-use traversal
 
-  // Remove outer uses that occur before the end of the borrow scope.
   auto *beginBorrow = cast<BeginBorrowInst>(currentDef);
-  BorrowedValue::get(beginBorrow).visitLocalScopeEndingUses([&](Operand *use) {
-    // Forward iterate until we find the end of the borrow scope.
-    auto *endScope = use->getUser();
-    for (auto instIter = beginBorrow->getIterator(),
-              endIter = endScope->getIterator();
-         instIter != endIter; ++instIter) {
-      outerUseInsts.erase(&*instIter);
-    }
-  });
+  SmallVector<SILInstruction *, 1> scopeEndingInst;
+  BorrowedValue::get(beginBorrow)
+      .getLocalScopeEndingInstructions(scopeEndingInst);
+  assert(scopeEndingInst.size() == 1 && "expected single-block borrow");
+  // Remove outer uses that occur before the end of the borrow scope by
+  // forward iterating from begin_borrow to end_borrow.
+  for (auto instIter = beginBorrow->getIterator(),
+            endIter = scopeEndingInst[0]->getIterator();
+       instIter != endIter; ++instIter) {
+    outerUseInsts.erase(&*instIter);
+  }
   if (outerUseInsts.empty()) {
-    return;
+    return true;
   }
-  // Rewrite the outer uses.
+  // Rewrite the outer uses and record lifetime-ending uses.
+  SmallVector<Operand *, 4> consumingUses;
+  SmallPtrSet<SILInstruction *, 4> unclaimedConsumingUsers;
   this->outerCopy = createOuterCopy(beginBorrow);
   for (Operand *use : outerUses) {
     if (!outerUseInsts.count(use->getUser())) {
-      continue;
+      // The immediate use is within this borrow scope.
+      BorrowingOperand borrowOper(use);
+      if (borrowOper.kind == BorrowingOperandKind::Invalid) {
+        continue;
+      }
+      // For sub-borrows also check that the scope-ending instructions are
+      // within the scope.
+      if (borrowOper.visitLocalEndScopeUses([&](Operand *endBorrow) {
+            return !outerUseInsts.count(endBorrow->getUser());
+          })) {
+        continue;
+      }
     }
     LLVM_DEBUG(llvm::dbgs() << "  Use of outer copy " << *use->getUser());
     use->set(outerCopy);
+    if (use->isLifetimeEnding()) {
+      consumingUses.push_back(use);
+      unclaimedConsumingUsers.insert(use->getUser());
+    }
+  }
+  // Insert a destroy on the outer copy's lifetime frontier, or claim an
+  // existing consume.
+  ValueLifetimeAnalysis lifetimeAnalysis(outerCopy, outerUseInsts);
+  ValueLifetimeAnalysis::Frontier frontier;
+  bool result = lifetimeAnalysis.computeFrontier(
+      frontier, ValueLifetimeAnalysis::DontModifyCFG, deBlocks);
+  assert(result);
+  while (!frontier.empty()) {
+    auto *insertPt = frontier.pop_back_val();
+    if (unclaimedConsumingUsers.erase(&*std::prev(insertPt->getIterator()))) {
+      continue;
+    }
+    SILBuilderWithScope(insertPt).createDestroyValue(insertPt->getLoc(),
+                                                     outerCopy);
   }
+  // Add copies for consuming users of outerCopy.
+  for (auto *use : consumingUses) {
+    // If the user is still in the unclaimedConsumingUsers set, then it does not
+    // end the outer copy's lifetime and therefore requires a copy. Only one
+    // operand can be claimed as ending the lifetime, so return its user to the
+    // unclaimedConsumingUsers set after skipping the first copy.
+    auto iterAndInserted = unclaimedConsumingUsers.insert(use->getUser());
+    if (!iterAndInserted.second) {
+      copyLiveUse(use);
+    }
+  }
+  return true;
 }
 
 //===----------------------------------------------------------------------===//
@@ -650,20 +745,6 @@ void CanonicalizeOSSALifetime::findOrInsertDestroys() {
 // MARK: Step 3. Rewrite copies and destroys
 //===----------------------------------------------------------------------===//
 
-/// The lifetime extends beyond given consuming use. Copy the value.
-static void copyLiveUse(Operand *use) {
-  SILInstruction *user = use->getUser();
-  SILBuilderWithScope B(user->getIterator());
-
-  auto loc = RegularLocation::getAutoGeneratedLocation(
-    user->getLoc().getSourceLoc());
-  auto *copy = B.createCopyValue(loc, use->get());
-  use->set(copy);
-
-  ++NumCopiesGenerated;
-  LLVM_DEBUG(llvm::dbgs() << "  Copying at last use " << *copy);
-}
-
 /// Revisit the def-use chain of currentDef. Mark unneeded original
 /// copies and destroys for deletion. Insert new copies for interior uses that
 /// require ownership of the used operand.
@@ -774,6 +855,8 @@ void CanonicalizeOSSALifetime::rewriteCopies() {
 //===----------------------------------------------------------------------===//
 
 bool CanonicalizeOSSALifetime::canonicalizeValueLifetime(SILValue def) {
+  LLVM_DEBUG(llvm::dbgs() << "  Canonicalizing: " << def);
+
   switch (def.getOwnershipKind()) {
   case OwnershipKind::None:
   case OwnershipKind::Unowned:
@@ -787,7 +870,10 @@ bool CanonicalizeOSSALifetime::canonicalizeValueLifetime(SILValue def) {
     }
     // Set outerCopy and persistentCopies and rewrite uses
     // outside the scope.
-    consolidateBorrowScope();
+    if (!consolidateBorrowScope()) {
+      clearLiveness();
+      return false;
+    }
     // Invalidate book-keeping before deleting instructions.
     clearLiveness();
     // Rewrite copies and delete extra destroys within the scope.