Merge pull request #35256 from gottesmm/pr-9c2c0f1fec0527092a68ca8fd955d6a5dc77799d

[sil-inst-opt] Change InstModCallbacks to specify a setUseValue callback instead of RAUW callbacks.

Author: gottesmm

include/swift/SILOptimizer/Utils/CanonicalizeInstruction.h

@@ -49,8 +49,8 @@ struct CanonicalizeInstruction {
         callbacks(
             [&](SILInstruction *toDelete) { killInstruction(toDelete); },
             [&](SILInstruction *newInst) { notifyNewInstruction(newInst); },
-            [&](SILValue oldValue, SILValue newValue) {
-              oldValue->replaceAllUsesWith(newValue);
+            [&](Operand *use, SILValue newValue) {
+              use->set(newValue);
               notifyHasNewUsers(newValue);
             }) {
 #ifndef NDEBUG

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -330,77 +330,75 @@ class InstModCallbacks {
   /// Default implementation is a no-op, but we still mark madeChange.
   std::function<void(SILInstruction *newlyCreatedInst)> createdNewInstFunc;
 
-  /// A function that first replaces all uses of oldValue with uses of newValue.
+  /// A function sets the value in \p use to be \p newValue.
   ///
-  /// Default implementation just calls oldValue->replaceAllUsesWith(newValue);
-  std::function<void(SILValue oldValue, SILValue newValue)>
-      replaceValueUsesWithFunc;
-
-  /// A function that first replaces all uses of oldValue with uses of newValue
-  /// and then erases oldValue.
-  ///
-  /// Default implementation just calls replaceValueUsesWithFunc and then
-  /// deleteInstFunc.
-  std::function<void(SingleValueInstruction *oldValue, SILValue newValue)>
-      eraseAndRAUWSingleValueInstFunc;
+  /// Default implementation just calls use->set(newValue).
+  std::function<void(Operand *use, SILValue newValue)> setUseValueFunc;
 
   /// A boolean that tracks if any of our callbacks were ever called.
-  bool madeChange = false;
+  bool wereAnyCallbacksInvoked = false;
 
 public:
   InstModCallbacks(decltype(deleteInstFunc) deleteInstFunc)
       : deleteInstFunc(deleteInstFunc) {}
 
   InstModCallbacks(decltype(deleteInstFunc) deleteInstFunc,
-                   decltype(createdNewInstFunc) createdNewInstFunc,
-                   decltype(replaceValueUsesWithFunc) replaceValueUsesWithFunc)
-      : deleteInstFunc(deleteInstFunc), createdNewInstFunc(createdNewInstFunc),
-        replaceValueUsesWithFunc(replaceValueUsesWithFunc) {}
+                   decltype(createdNewInstFunc) createdNewInstFunc)
+      : deleteInstFunc(deleteInstFunc), createdNewInstFunc(createdNewInstFunc) {
+  }
+
+  InstModCallbacks(decltype(deleteInstFunc) deleteInstFunc,
+                   decltype(setUseValueFunc) setUseValueFunc)
+      : deleteInstFunc(deleteInstFunc), setUseValueFunc(setUseValueFunc) {}
 
-  InstModCallbacks(
-      decltype(deleteInstFunc) deleteInstFunc,
-      decltype(createdNewInstFunc) createdNewInstFunc,
-      decltype(replaceValueUsesWithFunc) replaceValueUsesWithFunc,
-      decltype(eraseAndRAUWSingleValueInstFunc) eraseAndRAUWSingleValueInstFunc)
+  InstModCallbacks(decltype(deleteInstFunc) deleteInstFunc,
+                   decltype(createdNewInstFunc) createdNewInstFunc,
+                   decltype(setUseValueFunc) setUseValueFunc)
       : deleteInstFunc(deleteInstFunc), createdNewInstFunc(createdNewInstFunc),
-        replaceValueUsesWithFunc(replaceValueUsesWithFunc),
-        eraseAndRAUWSingleValueInstFunc(eraseAndRAUWSingleValueInstFunc) {}
+        setUseValueFunc(setUseValueFunc) {}
 
   InstModCallbacks() = default;
   ~InstModCallbacks() = default;
   InstModCallbacks(const InstModCallbacks &) = default;
   InstModCallbacks(InstModCallbacks &&) = default;
 
   void deleteInst(SILInstruction *instToDelete) {
-    madeChange = true;
+    wereAnyCallbacksInvoked = true;
     if (deleteInstFunc)
       return deleteInstFunc(instToDelete);
     instToDelete->eraseFromParent();
   }
 
   void createdNewInst(SILInstruction *newlyCreatedInst) {
-    madeChange = true;
+    wereAnyCallbacksInvoked = true;
     if (createdNewInstFunc)
       createdNewInstFunc(newlyCreatedInst);
   }
 
+  void setUseValue(Operand *use, SILValue newValue) {
+    wereAnyCallbacksInvoked = true;
+    if (setUseValueFunc)
+      return setUseValueFunc(use, newValue);
+    use->set(newValue);
+  }
+
   void replaceValueUsesWith(SILValue oldValue, SILValue newValue) {
-    madeChange = true;
-    if (replaceValueUsesWithFunc)
-      return replaceValueUsesWithFunc(oldValue, newValue);
-    oldValue->replaceAllUsesWith(newValue);
+    wereAnyCallbacksInvoked = true;
+
+    while (!oldValue->use_empty()) {
+      auto *use = *oldValue->use_begin();
+      setUseValue(use, newValue);
+    }
   }
 
   void eraseAndRAUWSingleValueInst(SingleValueInstruction *oldInst,
                                    SILValue newValue) {
-    madeChange = true;
-    if (eraseAndRAUWSingleValueInstFunc)
-      return eraseAndRAUWSingleValueInstFunc(oldInst, newValue);
+    wereAnyCallbacksInvoked = true;
     replaceValueUsesWith(oldInst, newValue);
     deleteInst(oldInst);
   }
 
-  bool getMadeChange() const { return madeChange; }
+  bool hadCallbackInvocation() const { return wereAnyCallbacksInvoked; }
 };
 
 /// Get all consumed arguments of a partial_apply.

lib/SILOptimizer/Analysis/SimplifyInstruction.cpp

@@ -752,7 +752,7 @@ swift::replaceAllUsesAndErase(SingleValueInstruction *svi, SILValue newValue,
       callbacks.deleteInst(user);
       continue;
     }
-    use->set(newValue);
+    callbacks.setUseValue(use, newValue);
   }
 
   callbacks.deleteInst(svi);

lib/SILOptimizer/Mandatory/MandatoryCombine.cpp

@@ -148,8 +148,9 @@ class MandatoryCombiner final
               instructionsPendingDeletion.push_back(instruction);
             },
             [&](SILInstruction *instruction) { worklist.add(instruction); },
-            [this](SILValue oldValue, SILValue newValue) {
-              worklist.replaceValueUsesWith(oldValue, newValue);
+            [this](Operand *use, SILValue newValue) {
+              use->set(newValue);
+              worklist.add(use->getUser());
             }),
         createdInstructions(createdInstructions),
         deadEndBlocks(deadEndBlocks){};

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -84,6 +84,9 @@ class SILCombiner :
   /// Cast optimizer
   CastOptimizer CastOpt;
 
+  /// Centralized InstModCallback that we use for certain utility methods.
+  InstModCallbacks instModCallbacks;
+
 public:
   SILCombiner(SILOptFunctionBuilder &FuncBuilder, SILBuilder &B,
               AliasAnalysis *AA, DominanceAnalysis *DA,
@@ -103,7 +106,18 @@ class SILCombiner :
               replaceInstUsesWith(*I, V);
             },
             /* EraseAction */
-            [&](SILInstruction *I) { eraseInstFromFunction(*I); }) {}
+            [&](SILInstruction *I) { eraseInstFromFunction(*I); }),
+        instModCallbacks(
+            [&](SILInstruction *instToDelete) {
+              eraseInstFromFunction(*instToDelete);
+            },
+            [&](SILInstruction *newlyCreatedInst) {
+              Worklist.add(newlyCreatedInst);
+            },
+            [&](Operand *use, SILValue newValue) {
+              use->set(newValue);
+              Worklist.add(use->getUser());
+            }) {}
 
   bool runOnFunction(SILFunction &F);
 
@@ -311,14 +325,7 @@ class SILCombiner :
                                        StringRef FInverseName, StringRef FName);
 
 private:
-  InstModCallbacks getInstModCallbacks() {
-    return InstModCallbacks(
-        [this](SILInstruction *DeadInst) { eraseInstFromFunction(*DeadInst); },
-        [this](SILInstruction *NewInst) { Worklist.add(NewInst); },
-        [this](SILValue oldValue, SILValue newValue) {
-          replaceValueUsesWith(oldValue, newValue);
-        });
-  }
+  InstModCallbacks &getInstModCallbacks() { return instModCallbacks; }
 
   // Build concrete existential information using findInitExistential.
   Optional<ConcreteOpenedExistentialInfo>

lib/SILOptimizer/SemanticARC/SemanticARCOptVisitor.h

@@ -53,9 +53,9 @@ struct LLVM_LIBRARY_VISIBILITY SemanticARCOptVisitor
       : ctx(fn, onlyGuaranteedOpts,
             InstModCallbacks(
                 [this](SILInstruction *inst) { eraseInstruction(inst); },
-                [](SILInstruction *) {}, [](SILValue, SILValue) {},
-                [this](SingleValueInstruction *i, SILValue value) {
-                  eraseAndRAUWSingleValueInstruction(i, value);
+                [this](Operand *use, SILValue newValue) {
+                  use->set(newValue);
+                  worklist.insert(newValue);
                 })) {}
 
   void reset() {
@@ -118,14 +118,7 @@ struct LLVM_LIBRARY_VISIBILITY SemanticARCOptVisitor
     drainVisitedSinceLastMutationIntoWorklist();
   }
 
-  InstModCallbacks getCallbacks() {
-    return InstModCallbacks(
-        [this](SILInstruction *inst) { eraseInstruction(inst); },
-        [](SILInstruction *) {}, [](SILValue, SILValue) {},
-        [this](SingleValueInstruction *i, SILValue value) {
-          eraseAndRAUWSingleValueInstruction(i, value);
-        });
-  }
+  InstModCallbacks &getCallbacks() { return ctx.instModCallbacks; }
 
   bool visitSILInstruction(SILInstruction *i) {
     assert((isa<OwnershipForwardingTermInst>(i) ||

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -434,7 +434,7 @@ void OwnershipRAUWUtility::eliminateReborrowsOfRecursiveBorrows(
 
     // Then set our borrowing operand to take our innerBorrow instead of value
     // (whose lifetime we just ended).
-    borrowingOperand->set(innerBorrow);
+    callbacks.setUseValue(*borrowingOperand, innerBorrow);
     // Add our outer end borrow as a use point to make sure that we extend our
     // base value to this point.
     usePoints.push_back(&outerEndBorrow->getAllOperands()[0]);
@@ -500,7 +500,7 @@ void OwnershipRAUWUtility::rewriteReborrows(
     callbacks.createdNewInst(innerBorrow);
     callbacks.createdNewInst(outerEndBorrow);
 
-    reborrow->set(innerBorrow);
+    callbacks.setUseValue(*reborrow, innerBorrow);
 
     auto *borrowedArg =
         const_cast<SILPhiArgument *>(bi->getArgForOperand(reborrow.op));
@@ -570,7 +570,7 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleUnowned() {
           auto *newInst = builder.createUncheckedOwnershipConversion(
               ti->getLoc(), use->get(), OwnershipKind::Unowned);
           callbacks.createdNewInst(newInst);
-          use->set(newInst);
+          callbacks.setUseValue(use, newInst);
         }
       }
     }
@@ -602,7 +602,7 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleUnowned() {
           auto *newInst = builder.createUncheckedOwnershipConversion(
               ti->getLoc(), use->get(), OwnershipKind::Unowned);
           callbacks.createdNewInst(newInst);
-          use->set(newInst);
+          callbacks.setUseValue(use, newInst);
         }
       }
     }