[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().

Author: vporpo

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 {

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();

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) {

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());
+}