[ownership] Change CanonicalOSSALifetime to use an InstModCallback.

gottesmm · gottesmm · commit c3c2b84e48e3 · 2021-04-21T22:09:55.000-07:00
This enables passes to use this as a utility that properly composes with how the
pass maintains its state. If an InstModCallback isn't passed in, we use the
default InstModCallback which should be cheap (always succeeding check for
nullptr + call inline default callback).
diff --git a/include/swift/SILOptimizer/Utils/CanonicalOSSALifetime.h b/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; }
 
diff --git a/include/swift/SILOptimizer/Utils/InstOptUtils.h b/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.
diff --git a/lib/SILOptimizer/Utils/CanonicalOSSALifetime.cpp b/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);
     };
