[llvm] Register callbacks in Region for instruction creation/deletion. (PR #117088)

[Jorge Gorbe Moya via llvm-commits]

https://github.com/slackito created https://github.com/llvm/llvm-project/pull/117088

This will keep the current Region updated when region passes add/delete instructions.

>From 9cca95e34e00641e1cf446a5ce8c7cf7ea6e3be0 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Wed, 20 Nov 2024 16:33:16 -0800
Subject: [PATCH] Register callbacks in Region for instruction
 creation/deletion.

This will keep the current Region updated when region passes add/delete
instructions.
---
 llvm/include/llvm/SandboxIR/Region.h    |  5 ++++
 llvm/lib/SandboxIR/Region.cpp           | 10 +++++++-
 llvm/unittests/SandboxIR/RegionTest.cpp | 31 +++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h
index 67411f3fb741da..8133e01734ea78 100644
--- a/llvm/include/llvm/SandboxIR/Region.h
+++ b/llvm/include/llvm/SandboxIR/Region.h
@@ -63,6 +63,11 @@ class Region {
 
   Context &Ctx;
 
+  /// ID (for later deregistration) of the "create instruction" callback.
+  Context::CallbackID CreateInstCB;
+  /// ID (for later deregistration) of the "erase instruction" callback.
+  Context::CallbackID EraseInstCB;
+
   // TODO: Add cost modeling.
   // TODO: Add a way to encode/decode region info to/from metadata.
 
diff --git a/llvm/lib/SandboxIR/Region.cpp b/llvm/lib/SandboxIR/Region.cpp
index b6292f3b24b876..1455012440f90c 100644
--- a/llvm/lib/SandboxIR/Region.cpp
+++ b/llvm/lib/SandboxIR/Region.cpp
@@ -15,9 +15,17 @@ Region::Region(Context &Ctx) : Ctx(Ctx) {
   LLVMContext &LLVMCtx = Ctx.LLVMCtx;
   auto *RegionStrMD = MDString::get(LLVMCtx, RegionStr);
   RegionMDN = MDNode::getDistinct(LLVMCtx, {RegionStrMD});
+
+  CreateInstCB = Ctx.registerCreateInstrCallback(
+      [this](Instruction *NewInst) { add(NewInst); });
+  EraseInstCB = Ctx.registerEraseInstrCallback(
+      [this](Instruction *ErasedInst) { remove(ErasedInst); });
 }
 
-Region::~Region() {}
+Region::~Region() {
+  Ctx.unregisterCreateInstrCallback(CreateInstCB);
+  Ctx.unregisterEraseInstrCallback(EraseInstCB);
+}
 
 void Region::add(Instruction *I) {
   Insts.insert(I);
diff --git a/llvm/unittests/SandboxIR/RegionTest.cpp b/llvm/unittests/SandboxIR/RegionTest.cpp
index a2efe551c8ff2c..47368f93a32c0c 100644
--- a/llvm/unittests/SandboxIR/RegionTest.cpp
+++ b/llvm/unittests/SandboxIR/RegionTest.cpp
@@ -81,6 +81,37 @@ define i8 @foo(i8 %v0, i8 %v1) {
 #endif
 }
 
+TEST_F(RegionTest, CallbackUpdates) {
+  parseIR(C, R"IR(
+define i8 @foo(i8 %v0, i8 %v1, ptr %ptr) {
+  %t0 = add i8 %v0, 1
+  %t1 = add i8 %t0, %v1
+  ret i8 %t0
+}
+)IR");
+  llvm::Function *LLVMF = &*M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto *F = Ctx.createFunction(LLVMF);
+  auto *Ptr = F->getArg(2);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  auto *T0 = cast<sandboxir::Instruction>(&*It++);
+  auto *T1 = cast<sandboxir::Instruction>(&*It++);
+  auto *Ret = cast<sandboxir::Instruction>(&*It++);
+  sandboxir::Region Rgn(Ctx);
+  Rgn.add(T0);
+  Rgn.add(T1);
+
+  // Test creation.
+  auto *NewI = sandboxir::StoreInst::create(T0, Ptr, /*Align=*/std::nullopt,
+                                            Ret->getIterator(), Ctx);
+  EXPECT_THAT(Rgn.insts(), testing::ElementsAre(T0, T1, NewI));
+
+  // Test deletion.
+  T1->eraseFromParent();
+  EXPECT_THAT(Rgn.insts(), testing::ElementsAre(T0, NewI));
+}
+
 TEST_F(RegionTest, MetadataFromIR) {
   parseIR(C, R"IR(
 define i8 @foo(i8 %v0, i8 %v1) {