1 files changed, 8 insertions, 5 deletions
diff --git a/llvm/lib/SandboxIR/User.cpp b/llvm/lib/SandboxIR/User.cpp
index d7e4656..43fd565 100644
--- a/llvm/lib/SandboxIR/User.cpp
+++ b/llvm/lib/SandboxIR/User.cpp
@@ -90,17 +90,20 @@ bool User::classof(const Value *From) {
 
 void User::setOperand(unsigned OperandIdx, Value *Operand) {
   assert(isa<llvm::User>(Val) && "No operands!");
-  Ctx.getTracker().emplaceIfTracking<UseSet>(getOperandUse(OperandIdx));
+  const auto &U = getOperandUse(OperandIdx);
+  Ctx.getTracker().emplaceIfTracking<UseSet>(U);
+  Ctx.runSetUseCallbacks(U, Operand);
   // We are delegating to llvm::User::setOperand().
   cast<llvm::User>(Val)->setOperand(OperandIdx, Operand->Val);
 }
 
 bool User::replaceUsesOfWith(Value *FromV, Value *ToV) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.isTracking()) {
-    for (auto OpIdx : seq<unsigned>(0, getNumOperands())) {
-      auto Use = getOperandUse(OpIdx);
-      if (Use.get() == FromV)
+  for (auto OpIdx : seq<unsigned>(0, getNumOperands())) {
+    auto Use = getOperandUse(OpIdx);
+    if (Use.get() == FromV) {
+      Ctx.runSetUseCallbacks(Use, ToV);
+      if (Tracker.isTracking())
         Tracker.emplaceIfTracking<UseSet>(Use);
     }
   }