[llvm] [SandboxIR][Tracker] Track creation of instructions (PR #102013)

[via llvm-commits]

https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/102013

>From a6937b6a441def5093454632e676c8f331f60a1e Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Fri, 19 Jul 2024 10:47:03 -0700
Subject: [PATCH] [SandboxIR][Tracker] Track creation of instructions

This patch adds tracking support for the creation of SandboxIR instructions.
This allows us to undo the creation of instructions during revert().

Please note that we don't allow the creation of detached instructions,
meaning that the instructions need to be inserted into a block upon creation.
This is why the tracker class combines creation and insertion.
---
 llvm/include/llvm/SandboxIR/Tracker.h    | 17 +++++++++++
 llvm/lib/SandboxIR/SandboxIR.cpp         | 10 +++++++
 llvm/lib/SandboxIR/Tracker.cpp           |  9 ++++++
 llvm/unittests/SandboxIR/TrackerTest.cpp | 36 ++++++++++++++++++++++++
 4 files changed, 72 insertions(+)

diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index f660ae94c8f7f..3c491bf0b3f58 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -328,6 +328,23 @@ class InsertIntoBB final : public IRChangeBase {
 #endif // NDEBUG
 };
 
+class CreateAndInsertInst final : public IRChangeBase {
+  Instruction *NewI = nullptr;
+
+public:
+  CreateAndInsertInst(Instruction *NewI, Tracker &Tracker)
+      : IRChangeBase(Tracker), NewI(NewI) {}
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "CreateAndInsertInst";
+  }
+  LLVM_DUMP_METHOD void dump() const final;
+#endif
+};
+
 /// The tracker collects all the change objects and implements the main API for
 /// saving / reverting / accepting.
 class Tracker {
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index 00abee48821ff..235c4574ead13 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -1359,6 +1359,16 @@ std::unique_ptr<Value> Context::detach(Value *V) {
 Value *Context::registerValue(std::unique_ptr<Value> &&VPtr) {
   assert(VPtr->getSubclassID() != Value::ClassID::User &&
          "Can't register a user!");
+
+  // Track creation of instructions.
+  // Please note that we don't allow the creation of detached instructions,
+  // meaning that the instructions need to be inserted into a block upon
+  // creation. This is why the tracker class combines creation and insertion.
+  auto &Tracker = getTracker();
+  if (Tracker.isTracking())
+    if (auto *I = dyn_cast<Instruction>(VPtr.get()))
+      Tracker.track(std::make_unique<CreateAndInsertInst>(I, Tracker));
+
   Value *V = VPtr.get();
   [[maybe_unused]] auto Pair =
       LLVMValueToValueMap.insert({VPtr->Val, std::move(VPtr)});
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 4ab24123cafc5..aa18c21cd03d6 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -300,6 +300,15 @@ void InsertIntoBB::dump() const {
 }
 #endif
 
+void CreateAndInsertInst::revert() { NewI->eraseFromParent(); }
+
+#ifndef NDEBUG
+void CreateAndInsertInst::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif
+
 void Tracker::track(std::unique_ptr<IRChangeBase> &&Change) {
   assert(State == TrackerState::Record && "The tracker should be tracking!");
   Changes.push_back(std::move(Change));
diff --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
index 76cee1105cbf0..e61258392a388 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -505,6 +505,42 @@ define void @foo(i32 %arg) {
   Add0->insertBefore(Ret);
 }
 
+// TODO: Test multi-instruction patterns.
+TEST_F(TrackerTest, CreateAndInsertInst) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr) {
+  %ld = load i8, ptr %ptr, align 64
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+
+  auto *F = Ctx.createFunction(&LLVMF);
+  auto *Ptr = F->getArg(0);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  auto *Ld = cast<sandboxir::LoadInst>(&*It++);
+  auto *Ret = &*It++;
+
+  Ctx.save();
+  // Check create(InsertBefore) with tracking enabled.
+  sandboxir::LoadInst *NewLd =
+      sandboxir::LoadInst::create(Ld->getType(), Ptr, Align(8),
+                                  /*InsertBefore=*/Ld, Ctx, "NewLd");
+  It = BB->begin();
+  EXPECT_EQ(&*It++, NewLd);
+  EXPECT_EQ(&*It++, Ld);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ld);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+}
+
 TEST_F(TrackerTest, CallBaseSetters) {
   parseIR(C, R"IR(
 declare void @bar1(i8)