[SandboxIR][Tracker] Track Instruction::moveBefore()

Track instruction movement. This supports both moveBefore() and moveAfter().

Author: vporpo

llvm/include/llvm/SandboxIR/Tracker.h

@@ -156,6 +156,26 @@ class RemoveFromParent : public IRChangeBase {
 #endif // NDEBUG
 };
 
+class MoveInstr : public IRChangeBase {
+  /// The instruction that moved.
+  Instruction *MovedI;
+  /// This is either the next instruction in the block, or the parent BB if at
+  /// the end of the BB.
+  PointerUnion<Instruction *, BasicBlock *> NextInstrOrBB;
+
+public:
+  MoveInstr(sandboxir::Instruction *I, Tracker &Tracker);
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "MoveInstr";
+  }
+  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

@@ -371,6 +371,11 @@ void Instruction::moveBefore(BasicBlock &BB, const BBIterator &WhereIt) {
   if (std::next(getIterator()) == WhereIt)
     // Destination is same as origin, nothing to do.
     return;
+
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<MoveInstr>(this, Tracker));
+
   auto *LLVMBB = cast<llvm::BasicBlock>(BB.Val);
   llvm::BasicBlock::iterator It;
   if (WhereIt == BB.end()) {

llvm/lib/SandboxIR/Tracker.cpp

@@ -124,6 +124,30 @@ void RemoveFromParent::dump() const {
 }
 #endif
 
+MoveInstr::MoveInstr(Instruction *MovedI, Tracker &Tracker)
+    : IRChangeBase(Tracker), MovedI(MovedI) {
+  if (auto *NextI = MovedI->getNextNode())
+    NextInstrOrBB = NextI;
+  else
+    NextInstrOrBB = MovedI->getParent();
+}
+
+void MoveInstr::revert() {
+  if (auto *NextI = NextInstrOrBB.dyn_cast<Instruction *>()) {
+    MovedI->moveBefore(NextI);
+  } else {
+    auto *BB = NextInstrOrBB.get<BasicBlock *>();
+    MovedI->moveBefore(*BB, BB->end());
+  }
+}
+
+#ifndef NDEBUG
+void MoveInstr::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));

llvm/unittests/SandboxIR/TrackerTest.cpp

@@ -255,3 +255,118 @@ define i32 @foo(i32 %arg) {
   EXPECT_EQ(&*It++, Ret);
   EXPECT_EQ(It, BB->end());
 }
+
+// TODO: Test multi-instruction patterns.
+TEST_F(TrackerTest, MoveInstr) {
+  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 *BB = &*F->begin();
+  auto It = BB->begin();
+  sandboxir::Instruction *Add0 = &*It++;
+  sandboxir::Instruction *Add1 = &*It++;
+  sandboxir::Instruction *Ret = &*It++;
+
+  // Check moveBefore(Instruction *) with tracking enabled.
+  Ctx.save();
+  Add1->moveBefore(Add0);
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // Same for the last instruction in the block.
+  Ctx.save();
+  Ret->moveBefore(Add0);
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(It, BB->end());
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // Check moveBefore(BasicBlock &, BasicBlock::iterator) with tracking enabled.
+  Ctx.save();
+  Add1->moveBefore(*BB, Add0->getIterator());
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // Same for the last instruction in the block.
+  Ctx.save();
+  Ret->moveBefore(*BB, Add0->getIterator());
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // Check moveAfter(Instruction *) with tracking enabled.
+  Ctx.save();
+  Add0->moveAfter(Add1);
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+
+  // Same for the last instruction in the block.
+  Ctx.save();
+  Ret->moveAfter(Add0);
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Add0);
+  EXPECT_EQ(&*It++, Add1);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+}