1 files changed, 18 insertions, 0 deletions
diff --git a/llvm/lib/SandboxIR/Context.cpp b/llvm/lib/SandboxIR/Context.cpp
index 830f283..6a397b0 100644
--- a/llvm/lib/SandboxIR/Context.cpp
+++ b/llvm/lib/SandboxIR/Context.cpp
@@ -687,6 +687,11 @@ void Context::runMoveInstrCallbacks(Instruction *I, const BBIterator &WhereIt) {
     CBEntry.second(I, WhereIt);
 }
 
+void Context::runSetUseCallbacks(const Use &U, Value *NewSrc) {
+  for (auto &CBEntry : SetUseCallbacks)
+    CBEntry.second(U, NewSrc);
+}
+
 // An arbitrary limit, to check for accidental misuse. We expect a small number
 // of callbacks to be registered at a time, but we can increase this number if
 // we discover we needed more.
@@ -732,4 +737,17 @@ void Context::unregisterMoveInstrCallback(CallbackID ID) {
          "Callback ID not found in MoveInstrCallbacks during deregistration");
 }
 
+Context::CallbackID Context::registerSetUseCallback(SetUseCallback CB) {
+  assert(SetUseCallbacks.size() <= MaxRegisteredCallbacks &&
+         "SetUseCallbacks size limit exceeded");
+  CallbackID ID{NextCallbackID++};
+  SetUseCallbacks[ID] = CB;
+  return ID;
+}
+void Context::unregisterSetUseCallback(CallbackID ID) {
+  [[maybe_unused]] bool Erased = SetUseCallbacks.erase(ID);
+  assert(Erased &&
+         "Callback ID not found in SetUseCallbacks during deregistration");
+}
+
 } // namespace llvm::sandboxir