[ownership] Change CanonicalOSSALifetime to use an InstModCallback.

include/swift/SILOptimizer/Utils/CanonicalOSSALifetime.h

@@ -97,6 +97,7 @@
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SILOptimizer/Analysis/NonLocalAccessBlockAnalysis.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/PrunedLiveness.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
@@ -107,13 +108,19 @@ namespace swift {
 /// result or invalid SILValue. The caller must delete the extract and its
 /// now-dead copy use.
 ///
-// If a copied-def is a struct-extract, attempt a destructure conversion
-//   %extract = struct_extract %... : $TypeWithSingleOwnershipValue
-//   %copy = copy_value %extract : $OwnershipValue
-// To:
-//   %copy = copy_value %extract : $TypeWithSingleOwnershipValue
-//   (%extracted,...) = destructure %copy : $TypeWithSingleOwnershipValue
-SILValue convertExtractToDestructure(StructExtractInst *extract);
+/// If a copied-def is a struct-extract, attempt a destructure conversion
+///   %extract = struct_extract %... : $TypeWithSingleOwnershipValue
+///   %copy = copy_value %extract : $OwnershipValue
+/// To:
+///   %copy = copy_value %extract : $TypeWithSingleOwnershipValue
+///   (%extracted,...) = destructure %copy : $TypeWithSingleOwnershipValue
+///
+/// \p instModCallbacks If non-null, this routine uses
+/// InstModCallbacks::{setUseValue,RAUW}() internally to modify code. Otherwise,
+/// just performs standard operations.
+SILValue
+convertExtractToDestructure(StructExtractInst *extract,
+                            InstModCallbacks *instModCallbacks = nullptr);
 
 /// Information about consumes on the extended-lifetime boundary. Consuming uses
 /// within the lifetime are not included--they will consume a copy after
@@ -333,16 +340,21 @@ class CanonicalizeOSSALifetime {
   /// lifetime.
   CanonicalOSSAConsumeInfo consumes;
 
+  /// The callbacks to use when deleting/rauwing instructions.
+  InstModCallbacks instModCallbacks;
+
 public:
-  CanonicalizeOSSALifetime(bool pruneDebugMode, bool canonicalizeBorrowMode,
-                           bool poisonRefsMode,
-                           NonLocalAccessBlockAnalysis *accessBlockAnalysis,
-                           DominanceAnalysis *dominanceAnalysis)
-    : pruneDebugMode(pruneDebugMode),
-      canonicalizeBorrowMode(canonicalizeBorrowMode),
-      poisonRefsMode(poisonRefsMode),
-      accessBlockAnalysis(accessBlockAnalysis),
-      dominanceAnalysis(dominanceAnalysis) {}
+  CanonicalizeOSSALifetime(
+      bool pruneDebugMode, bool canonicalizeBorrowMode, bool poisonRefsMode,
+      NonLocalAccessBlockAnalysis *accessBlockAnalysis,
+      DominanceAnalysis *dominanceAnalysis,
+      InstModCallbacks instModCallbacks = InstModCallbacks())
+      : pruneDebugMode(pruneDebugMode),
+        canonicalizeBorrowMode(canonicalizeBorrowMode),
+        poisonRefsMode(poisonRefsMode),
+        accessBlockAnalysis(accessBlockAnalysis),
+        dominanceAnalysis(dominanceAnalysis),
+        instModCallbacks(instModCallbacks) {}
 
   SILValue getCurrentDef() const { return currentDef; }
 

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -361,6 +361,9 @@ class InstModCallbacks {
   /// A function sets the value in \p use to be \p newValue.
   ///
   /// Default implementation just calls use->set(newValue).
+  ///
+  /// NOTE: It is assumed that this operation will never invalidate instruction
+  /// iterators.
   std::function<void(Operand *use, SILValue newValue)> setUseValueFunc;
 
   /// A boolean that tracks if any of our callbacks were ever called.

lib/SILOptimizer/Utils/CanonicalOSSALifetime.cpp

@@ -94,20 +94,22 @@ SILValue CanonicalizeOSSALifetime::getCanonicalCopiedDef(SILValue v) {
 /// The lifetime extends beyond given consuming use. Copy the value.
 ///
 /// This can set the operand value, but cannot invalidate the use iterator.
-static void copyLiveUse(Operand *use) {
+static void copyLiveUse(Operand *use, InstModCallbacks &instModCallbacks) {
   SILInstruction *user = use->getUser();
-  SILBuilderWithScope B(user->getIterator());
+  SILBuilderWithScope builder(user->getIterator());
 
   auto loc = RegularLocation::getAutoGeneratedLocation(user->getLoc());
-  auto *copy = B.createCopyValue(loc, use->get());
-  use->set(copy);
+  auto *copy = builder.createCopyValue(loc, use->get());
+  instModCallbacks.createdNewInst(copy);
+  instModCallbacks.setUseValue(use, copy);
 
   ++NumCopiesGenerated;
   LLVM_DEBUG(llvm::dbgs() << "  Copying at last use " << *copy);
 }
 
 // TODO: generalize this to handle multiple nondebug uses of the struct_extract.
-SILValue swift::convertExtractToDestructure(StructExtractInst *extract) {
+SILValue swift::convertExtractToDestructure(StructExtractInst *extract,
+                                            InstModCallbacks *callbacks) {
   if (!hasOneNonDebugUse(extract))
     return nullptr;
 
@@ -123,12 +125,20 @@ SILValue swift::convertExtractToDestructure(StructExtractInst *extract) {
   auto loc = extract->getLoc();
   auto *copy = builder.createCopyValue(loc, extract->getOperand());
   auto *destructure = builder.createDestructureStruct(loc, copy);
+  if (callbacks) {
+    callbacks->createdNewInst(copy);
+    callbacks->createdNewInst(destructure);
+  }
 
   SILValue nonTrivialResult = destructure->getResult(extract->getFieldIndex());
   assert(!nonTrivialResult->getType().isTrivial(*destructure->getFunction())
          && "field idx mismatch");
 
-  extractCopy->replaceAllUsesWith(nonTrivialResult);
+  if (callbacks) {
+    callbacks->replaceValueUsesWith(extractCopy, nonTrivialResult);
+  } else {
+    extractCopy->replaceAllUsesWith(nonTrivialResult);
+  }
   return nonTrivialResult;
 }
 
@@ -176,11 +186,13 @@ bool CanonicalizeOSSALifetime::computeBorrowLiveness() {
 // To use an existing outer copy, we could find its earliest consume. But the
 // new copy will immediately canonicalized and a canonical begin_borrow scope
 // have no outer uses of its first block.
-static CopyValueInst *createOuterCopy(BeginBorrowInst *beginBorrow) {
-  SILBuilderWithScope B(beginBorrow);
+static CopyValueInst *createOuterCopy(BeginBorrowInst *beginBorrow,
+                                      InstModCallbacks &instModCallbacks) {
+  SILBuilderWithScope builder(beginBorrow);
 
   auto loc = RegularLocation::getAutoGeneratedLocation(beginBorrow->getLoc());
-  auto *copy = B.createCopyValue(loc, beginBorrow->getOperand());
+  auto *copy = builder.createCopyValue(loc, beginBorrow->getOperand());
+  instModCallbacks.createdNewInst(copy);
 
   ++NumCopiesGenerated;
   LLVM_DEBUG(llvm::dbgs() << "  Outer copy " << *copy);
@@ -346,7 +358,7 @@ void CanonicalizeOSSALifetime::rewriteOuterBorrowUsesAndFindConsumes(
 
   auto rewriteOuterUse = [&](Operand *use) {
     LLVM_DEBUG(llvm::dbgs() << "  Use of outer copy " << *use->getUser());
-    use->set(newIncomingValue);
+    instModCallbacks.setUseValue(use, newIncomingValue);
     currentOuterUseInsts.push_back(use->getUser());
     outerUseInsts.insert(incomingValue->getDefiningInstruction());
     if (use->isLifetimeEnding()) {
@@ -444,18 +456,19 @@ void CanonicalizeOSSALifetime::rewriteOuterBorrowUsesAndFindConsumes(
   for (auto *boundaryEdge : boundary.boundaryEdges) {
     if (DeadEndBlocks::triviallyEndsInUnreachable(boundaryEdge))
       continue;
-
     auto insertPt = boundaryEdge->begin();
-    auto loc = insertPt->getLoc();
-    SILBuilderWithScope(insertPt).createDestroyValue(loc, newIncomingValue);
+    auto *dvi = SILBuilderWithScope(insertPt).createDestroyValue(
+        insertPt->getLoc(), newIncomingValue);
+    instModCallbacks.createdNewInst(dvi);
   }
 
   for (SILInstruction *lastUser : boundary.lastUsers) {
     if (unclaimedConsumingUsers.erase(lastUser))
         continue;
 
     SILBuilderWithScope::insertAfter(lastUser, [&](SILBuilder &b) {
-      b.createDestroyValue(lastUser->getLoc(), newIncomingValue);
+      auto *dvi = b.createDestroyValue(lastUser->getLoc(), newIncomingValue);
+      instModCallbacks.createdNewInst(dvi);
     });
   }
   // Add copies for consuming users of newIncomingValue.
@@ -466,7 +479,7 @@ void CanonicalizeOSSALifetime::rewriteOuterBorrowUsesAndFindConsumes(
     // unclaimedConsumingUsers set after skipping the first copy.
     auto iterAndInserted = unclaimedConsumingUsers.insert(use->getUser());
     if (!iterAndInserted.second) {
-      copyLiveUse(use);
+      copyLiveUse(use, instModCallbacks);
     }
   }
 }
@@ -496,7 +509,8 @@ bool CanonicalizeOSSALifetime::consolidateBorrowScope() {
   if (outerUseInsts.empty()) {
     return true;
   }
-  this->outerCopy = createOuterCopy(cast<BeginBorrowInst>(currentDef));
+  this->outerCopy =
+      createOuterCopy(cast<BeginBorrowInst>(currentDef), instModCallbacks);
 
   defUseWorklist.clear();
   rewriteOuterBorrowUsesAndFindConsumes(currentDef, outerUseInsts);
@@ -761,6 +775,7 @@ void CanonicalizeOSSALifetime::insertDestroyOnCFGEdge(
   SILBuilderWithScope builder(pos);
   auto loc = RegularLocation::getAutoGeneratedLocation(pos->getLoc());
   auto *di = builder.createDestroyValue(loc, currentDef, needsPoison);
+  instModCallbacks.createdNewInst(di);
 
   consumes.recordFinalConsume(di);
 
@@ -772,9 +787,10 @@ void CanonicalizeOSSALifetime::insertDestroyOnCFGEdge(
 ///
 /// Create a final destroy, immediately after `pos`.
 static void insertDestroyAtInst(SILBasicBlock::iterator pos,
-                                DestroyValueInst *existingDestroy,
-                                SILValue def, bool needsPoison,
-                                CanonicalOSSAConsumeInfo &consumes) {
+                                DestroyValueInst *existingDestroy, SILValue def,
+                                bool needsPoison,
+                                CanonicalOSSAConsumeInfo &consumes,
+                                InstModCallbacks &callbacks) {
   if (existingDestroy) {
     for (; pos != existingDestroy->getIterator(); ++pos) {
       if (auto *debugVal = dyn_cast<DebugValueInst>(&*pos)) {
@@ -791,6 +807,7 @@ static void insertDestroyAtInst(SILBasicBlock::iterator pos,
   SILBuilderWithScope builder(pos);
   auto loc = RegularLocation::getAutoGeneratedLocation((*pos).getLoc());
   auto *di = builder.createDestroyValue(loc, def, needsPoison);
+  callbacks.createdNewInst(di);
   consumes.recordFinalConsume(di);
 
   ++NumDestroysGenerated;
@@ -837,7 +854,7 @@ void CanonicalizeOSSALifetime::findOrInsertDestroyInBlock(SILBasicBlock *bb) {
           needsPoison = existingDestroyNeedsPoison;
         }
         insertDestroyAtInst(std::next(instIter), existingDestroy, currentDef,
-                            needsPoison, consumes);
+                            needsPoison, consumes, instModCallbacks);
         setChanged();
       }
       return;
@@ -889,7 +906,7 @@ void CanonicalizeOSSALifetime::findOrInsertDestroyInBlock(SILBasicBlock *bb) {
         needsPoison = existingDestroyNeedsPoison;
       }
       insertDestroyAtInst(instIter, existingDestroy, currentDef, needsPoison,
-                          consumes);
+                          consumes, instModCallbacks);
       setChanged();
       return;
     }
@@ -901,7 +918,7 @@ void CanonicalizeOSSALifetime::findOrInsertDestroyInBlock(SILBasicBlock *bb) {
         needsPoison = existingDestroyNeedsPoison;
       }
       insertDestroyAtInst(std::next(instIter), existingDestroy, currentDef,
-                          needsPoison, consumes);
+                          needsPoison, consumes, instModCallbacks);
       setChanged();
       return;
     }
@@ -1033,7 +1050,7 @@ void CanonicalizeOSSALifetime::rewriteCopies() {
       continue;
     }
     if (!visitUse(use)) {
-      copyLiveUse(use);
+      copyLiveUse(use, instModCallbacks);
       setChanged();
     }
   }
@@ -1047,16 +1064,16 @@ void CanonicalizeOSSALifetime::rewriteCopies() {
         if (!reusedCopyOp && srcCopy->getParent() == use->getParentBlock()) {
           reusedCopyOp = use;
         } else {
-          copyLiveUse(use);
+          copyLiveUse(use, instModCallbacks);
           setChanged();
         }
       }
     }
     if (!(reusedCopyOp && srcCopy->hasOneUse())) {
       setChanged();
-      srcCopy->replaceAllUsesWith(srcCopy->getOperand());
+      instModCallbacks.replaceValueUsesWith(srcCopy, srcCopy->getOperand());
       if (reusedCopyOp) {
-        reusedCopyOp->set(srcCopy);
+        instModCallbacks.setUseValue(reusedCopyOp, srcCopy);
       } else {
         if (instsToDelete.insert(srcCopy)) {
           LLVM_DEBUG(llvm::dbgs() << "  Removing " << *srcCopy);
@@ -1081,7 +1098,7 @@ void CanonicalizeOSSALifetime::rewriteCopies() {
   // Remove any dead, non-recovered debug_values.
   for (auto *dvi : consumes.getDebugInstsAfterConsume()) {
     LLVM_DEBUG(llvm::dbgs() << "  Removing debug_value: " << *dvi);
-    dvi->eraseFromParent();
+    instModCallbacks.deleteInst(dvi);
   }
 
   // Remove the leftover copy_value and destroy_value instructions.
@@ -1100,7 +1117,7 @@ void CanonicalizeOSSALifetime::injectPoison() {
         builder.getInsertionPoint()->getLoc());
       auto *di = builder.createDestroyValue(loc, currentDef,
                                             /*needsPoison*/ true);
-
+      instModCallbacks.createdNewInst(di);
       ++NumDestroysGenerated;
       LLVM_DEBUG(llvm::dbgs() << "  Destroy at last use " << *di);
     };