[llvm] 4f786c6 - [SandboxIR][Tracker] Track Instruction::removeFromParent() (#99541)

[via llvm-commits]

Author: vporpo
Date: 2024-07-18T13:54:54-07:00
New Revision: 4f786c6823956f72ef0b9eee57bf0c462842b609

URL: https://github.com/llvm/llvm-project/commit/4f786c6823956f72ef0b9eee57bf0c462842b609
DIFF: https://github.com/llvm/llvm-project/commit/4f786c6823956f72ef0b9eee57bf0c462842b609.diff

LOG: [SandboxIR][Tracker] Track Instruction::removeFromParent() (#99541)

This patch adds the necessary functionality to the Tracker and to the
Sandbox IR API functions for tracking calls to removeFromParent().

Added: 
    

Modified: 
    llvm/include/llvm/SandboxIR/Tracker.h
    llvm/lib/SandboxIR/SandboxIR.cpp
    llvm/lib/SandboxIR/Tracker.cpp
    llvm/unittests/SandboxIR/TrackerTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index a91b9f178b8aa..6332b5266bb31 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -53,6 +53,7 @@
 namespace llvm::sandboxir {
 
 class BasicBlock;
+class Instruction;
 class Tracker;
 
 /// The base class for IR Change classes.
@@ -135,6 +136,26 @@ class EraseFromParent : public IRChangeBase {
 #endif
 };
 
+class RemoveFromParent : public IRChangeBase {
+  /// The instruction that is about to get removed.
+  Instruction *RemovedI = nullptr;
+  /// This is either the next instr, or the parent BB if at the end of the BB.
+  PointerUnion<Instruction *, BasicBlock *> NextInstrOrBB;
+
+public:
+  RemoveFromParent(Instruction *RemovedI, Tracker &Tracker);
+  void revert() final;
+  void accept() final {};
+  Instruction *getInstruction() const { return RemovedI; }
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "RemoveFromParent";
+  }
+  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 c6daf1a586546..249ebd7366204 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -333,6 +333,10 @@ Instruction *Instruction::getPrevNode() const {
 }
 
 void Instruction::removeFromParent() {
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<RemoveFromParent>(this, Tracker));
+
   // Detach all the LLVM IR instructions from their parent BB.
   for (llvm::Instruction *I : getLLVMInstrs())
     I->removeFromParent();

diff  --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 2336a0067abbd..08ddb7b5c64c9 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -98,6 +98,30 @@ void EraseFromParent::dump() const {
   dump(dbgs());
   dbgs() << "\n";
 }
+#endif // NDEBUG
+
+RemoveFromParent::RemoveFromParent(Instruction *RemovedI, Tracker &Tracker)
+    : IRChangeBase(Tracker), RemovedI(RemovedI) {
+  if (auto *NextI = RemovedI->getNextNode())
+    NextInstrOrBB = NextI;
+  else
+    NextInstrOrBB = RemovedI->getParent();
+}
+
+void RemoveFromParent::revert() {
+  if (auto *NextI = NextInstrOrBB.dyn_cast<Instruction *>()) {
+    RemovedI->insertBefore(NextI);
+  } else {
+    auto *BB = NextInstrOrBB.get<BasicBlock *>();
+    RemovedI->insertInto(BB, BB->end());
+  }
+}
+
+#ifndef NDEBUG
+void RemoveFromParent::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
 #endif
 
 void Tracker::track(std::unique_ptr<IRChangeBase> &&Change) {

diff  --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
index 4e5dccf3a231f..e5c558345e187 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -198,3 +198,60 @@ define void @foo(i32 %v1) {
   EXPECT_EQ(&*It++, Ret);
   EXPECT_EQ(It, BB->end());
 }
+
+// TODO: Test multi-instruction patterns.
+TEST_F(TrackerTest, RemoveFromParent) {
+  parseIR(C, R"IR(
+define i32 @foo(i32 %arg) {
+  %add0 = add i32 %arg, %arg
+  %add1 = add i32 %add0, %arg
+  ret i32 %add1
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+
+  auto *F = Ctx.createFunction(&LLVMF);
+  auto *Arg = F->getArg(0);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  sandboxir::Instruction *Add0 = &*It++;
+  sandboxir::Instruction *Add1 = &*It++;
+  sandboxir::Instruction *Ret = &*It++;
+
+  Ctx.save();
+  // Check removeFromParent().
+  Add1->removeFromParent();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Removed instruction still be connected to operands and users.
+  EXPECT_EQ(Add1->getOperand(0), Add0);
+  EXPECT_EQ(Add1->getOperand(1), Arg);
+  EXPECT_EQ(Add0->getNumUses(), 1u);
+
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  EXPECT_EQ(Add1->getOperand(0), Add0);
+
+  // Same for the last instruction in the block.
+  Ctx.save();
+  Ret->removeFromParent();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(It, BB->end());
+  EXPECT_EQ(Ret->getOperand(0), Add1);
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+}