[SandboxIR] Add tracking for ShuffleVectorInst::setShuffleMask. (#105590)

Author: slackito

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -1024,9 +1024,7 @@ class ShuffleVectorInst final
   static Constant *convertShuffleMaskForBitcode(ArrayRef<int> Mask,
                                                 Type *ResultTy, Context &Ctx);
 
-  void setShuffleMask(ArrayRef<int> Mask) {
-    cast<llvm::ShuffleVectorInst>(Val)->setShuffleMask(Mask);
-  }
+  void setShuffleMask(ArrayRef<int> Mask);
 
   ArrayRef<int> getShuffleMask() const {
     return cast<llvm::ShuffleVectorInst>(Val)->getShuffleMask();

llvm/include/llvm/SandboxIR/Tracker.h

@@ -62,6 +62,7 @@ class AllocaInst;
 class CatchSwitchInst;
 class SwitchInst;
 class ConstantInt;
+class ShuffleVectorInst;
 
 /// The base class for IR Change classes.
 class IRChangeBase {
@@ -355,6 +356,20 @@ class CreateAndInsertInst final : public IRChangeBase {
 #endif
 };
 
+class ShuffleVectorSetMask final : public IRChangeBase {
+  ShuffleVectorInst *SVI;
+  SmallVector<int, 8> PrevMask;
+
+public:
+  ShuffleVectorSetMask(ShuffleVectorInst *SVI);
+  void revert(Tracker &Tracker) final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final { OS << "ShuffleVectorSetMask"; }
+  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 {

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -1868,6 +1868,11 @@ Value *ShuffleVectorInst::create(Value *V1, Value *V2, ArrayRef<int> Mask,
   return Ctx.getOrCreateConstant(cast<llvm::Constant>(NewV));
 }
 
+void ShuffleVectorInst::setShuffleMask(ArrayRef<int> Mask) {
+  Ctx.getTracker().emplaceIfTracking<ShuffleVectorSetMask>(this);
+  cast<llvm::ShuffleVectorInst>(Val)->setShuffleMask(Mask);
+}
+
 Constant *ShuffleVectorInst::getShuffleMaskForBitcode() const {
   return Ctx.getOrCreateConstant(
       cast<llvm::ShuffleVectorInst>(Val)->getShuffleMaskForBitcode());

llvm/lib/SandboxIR/Tracker.cpp

@@ -234,6 +234,20 @@ void CreateAndInsertInst::dump() const {
 }
 #endif
 
+ShuffleVectorSetMask::ShuffleVectorSetMask(ShuffleVectorInst *SVI)
+    : SVI(SVI), PrevMask(SVI->getShuffleMask()) {}
+
+void ShuffleVectorSetMask::revert(Tracker &Tracker) {
+  SVI->setShuffleMask(PrevMask);
+}
+
+#ifndef NDEBUG
+void ShuffleVectorSetMask::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif
+
 void Tracker::save() { State = TrackerState::Record; }
 
 void Tracker::revert() {

llvm/unittests/SandboxIR/TrackerTest.cpp

@@ -13,6 +13,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/SandboxIR/SandboxIR.h"
 #include "llvm/Support/SourceMgr.h"
+#include "gmock/gmock-matchers.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -792,6 +793,31 @@ define void @foo(i32 %cond0, i32 %cond1) {
   EXPECT_EQ(Switch->findCaseDest(BB1), One);
 }
 
+TEST_F(TrackerTest, ShuffleVectorInstSetters) {
+  parseIR(C, R"IR(
+define void @foo(<2 x i8> %v1, <2 x i8> %v2) {
+  %shuf = shufflevector <2 x i8> %v1, <2 x i8> %v2, <2 x i32> <i32 1, i32 2>
+  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();
+  auto *SVI = cast<sandboxir::ShuffleVectorInst>(&*It++);
+
+  // Check setShuffleMask.
+  SmallVector<int, 2> OrigMask(SVI->getShuffleMask());
+  Ctx.save();
+  SVI->setShuffleMask(ArrayRef<int>({0, 0}));
+  EXPECT_THAT(SVI->getShuffleMask(),
+              testing::Not(testing::ElementsAreArray(OrigMask)));
+  Ctx.revert();
+  EXPECT_THAT(SVI->getShuffleMask(), testing::ElementsAreArray(OrigMask));
+}
+
 TEST_F(TrackerTest, AtomicRMWSetters) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr, i8 %arg) {