From 1dbbabf95832d8bc89a0b474ae7d6844ffb0e18b Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas@google.com>
Date: Thu, 18 Jul 2024 15:10:15 -0700
Subject: [PATCH] [SandboxIR][Tracker] Track InsertIntoBB

This tracks the insertion of an Instruction into a BasicBlock.
---
 llvm/include/llvm/SandboxIR/Tracker.h    | 16 ++++++
 llvm/lib/SandboxIR/SandboxIR.cpp         | 14 +++++-
 llvm/lib/SandboxIR/Tracker.cpp           | 12 +++++
 llvm/unittests/SandboxIR/TrackerTest.cpp | 63 ++++++++++++++++++++++++
 4 files changed, 104 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index 238e4e9dacd34..8a963a41449ef 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -291,6 +291,22 @@ class MoveInstr : public IRChangeBase {
 #endif // NDEBUG
 };
 
+class InsertIntoBB final : public IRChangeBase {
+  Instruction *InsertedI = nullptr;
+
+public:
+  InsertIntoBB(Instruction *InsertedI, Tracker &Tracker);
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "InsertIntoBB";
+  }
+  LLVM_DUMP_METHOD void dump() const final;
+#endif // NDEBUG
+};
+
 /// 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 a67cb2a5bb25a..2f94c83311776 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -430,6 +430,11 @@ void Instruction::insertBefore(Instruction *BeforeI) {
   assert(is_sorted(getLLVMInstrs(),
                    [](auto *I1, auto *I2) { return I1->comesBefore(I2); }) &&
          "Expected program order!");
+
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<InsertIntoBB>(this, Tracker));
+
   // Insert the LLVM IR Instructions in program order.
   for (llvm::Instruction *I : getLLVMInstrs())
     I->insertBefore(BeforeTopI);
@@ -443,14 +448,21 @@ void Instruction::insertInto(BasicBlock *BB, const BBIterator &WhereIt) {
   llvm::BasicBlock *LLVMBB = cast<llvm::BasicBlock>(BB->Val);
   llvm::Instruction *LLVMBeforeI;
   llvm::BasicBlock::iterator LLVMBeforeIt;
+  Instruction *BeforeI;
   if (WhereIt != BB->end()) {
-    Instruction *BeforeI = &*WhereIt;
+    BeforeI = &*WhereIt;
     LLVMBeforeI = BeforeI->getTopmostLLVMInstruction();
     LLVMBeforeIt = LLVMBeforeI->getIterator();
   } else {
+    BeforeI = nullptr;
     LLVMBeforeI = nullptr;
     LLVMBeforeIt = LLVMBB->end();
   }
+
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<InsertIntoBB>(this, Tracker));
+
   // Insert the LLVM IR Instructions in program order.
   for (llvm::Instruction *I : getLLVMInstrs())
     I->insertInto(LLVMBB, LLVMBeforeIt);
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 0310160e8bf35..1f154d7de8948 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -259,6 +259,18 @@ void MoveInstr::dump() const {
 }
 #endif
 
+void InsertIntoBB::revert() { InsertedI->removeFromParent(); }
+
+InsertIntoBB::InsertIntoBB(Instruction *InsertedI, Tracker &Tracker)
+    : IRChangeBase(Tracker), InsertedI(InsertedI) {}
+
+#ifndef NDEBUG
+void InsertIntoBB::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 d016c7793a52c..04979b45755e0 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -442,6 +442,69 @@ define i32 @foo(i32 %arg) {
   EXPECT_EQ(It, BB->end());
 }
 
+// TODO: Test multi-instruction patterns.
+TEST_F(TrackerTest, InsertIntoBB) {
+  parseIR(C, R"IR(
+define void @foo(i32 %arg) {
+  %add0 = add i32 %arg, %arg
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+
+  auto *F = Ctx.createFunction(&LLVMF);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  sandboxir::Instruction *Add0 = &*It++;
+  sandboxir::Instruction *Ret = &*It++;
+  // Detach `Add0` before we save.
+  Add0->removeFromParent();
+
+  // Check insertBefore(Instruction *) with tracking enabled.
+  Ctx.save();
+  Add0->insertBefore(Ret);
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // Check insertAfter(Instruction *) with tracking enabled.
+  Ctx.save();
+  Add0->insertAfter(Ret);
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // Check insertInto(BasicBlock *, BasicBlock::iterator) with tracking enabled.
+  Ctx.save();
+  Add0->insertInto(BB, Ret->getIterator());
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // To make sure we don't leak memory insert `Add0` back into the BB before the
+  // end of the test.
+  Add0->insertBefore(Ret);
+}
+
 TEST_F(TrackerTest, CallBaseSetters) {
   parseIR(C, R"IR(
 declare void @bar1(i8)