[SandboxIR][NFC] GenericSetter tracker class (llvm#102260)

vporpo · web-flow · commit 4fe33d067c5d · 2024-08-07T17:17:19.000-07:00
This patch introduces the `GenericSetter` tracker class that can be used
to track and revert simple instruction setters.

This patch also replaces several setter tracker classes with the generic
one.
diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -239,71 +239,45 @@ class RemoveFromParent : public IRChangeBase {
 #endif // NDEBUG
 };
 
-class CallBrInstSetDefaultDest : public IRChangeBase {
-  CallBrInst *CallBr;
-  BasicBlock *OrigDefaultDest;
-
-public:
-  CallBrInstSetDefaultDest(CallBrInst *CallBr, Tracker &Tracker);
-  void revert() final;
-  void accept() final {}
-#ifndef NDEBUG
-  void dump(raw_ostream &OS) const final {
-    dumpCommon(OS);
-    OS << "CallBrInstSetDefaultDest";
-  }
-  LLVM_DUMP_METHOD void dump() const final;
-#endif
-};
-
-class AllocaSetAllocatedType final : public IRChangeBase {
-  AllocaInst *Alloca;
-  Type *OrigType;
-
-public:
-  AllocaSetAllocatedType(AllocaInst *Alloca, Tracker &Tracker);
-  void revert() final;
-  void accept() final {}
-#ifndef NDEBUG
-  void dump(raw_ostream &OS) const final {
-    dumpCommon(OS);
-    OS << "AllocaSetAllocatedType";
-  }
-  LLVM_DUMP_METHOD void dump() const final;
-#endif
-};
-
-class AllocaSetAlignment final : public IRChangeBase {
-  AllocaInst *Alloca;
-  Align OrigAlign;
+/// This class can be used for tracking most instruction setters.
+/// The two template arguments are:
+/// - GetterFn: The getter member function pointer (e.g., `&Foo::get`)
+/// - SetterFn: The setter member function pointer (e.g., `&Foo::set`)
+/// Upon construction, it saves a copy of the original value by calling the
+/// getter function. Revert sets the value back to the one saved, using the
+/// setter function provided.
+///
+/// Example:
+///  Tracker.track(std::make_unique<
+///                GenericSetter<&FooInst::get, &FooInst::set>>(I, Tracker));
+///
+template <auto GetterFn, auto SetterFn>
+class GenericSetter final : public IRChangeBase {
+  /// Helper for getting the class type from the getter
+  template <typename ClassT, typename RetT>
+  static ClassT getClassTypeFromGetter(RetT (ClassT::*Fn)() const);
+  template <typename ClassT, typename RetT>
+  static ClassT getClassTypeFromGetter(RetT (ClassT::*Fn)());
+
+  using InstrT = decltype(getClassTypeFromGetter(GetterFn));
+  using SavedValT = std::invoke_result_t<decltype(GetterFn), InstrT>;
+  InstrT *I;
+  SavedValT OrigVal;
 
 public:
-  AllocaSetAlignment(AllocaInst *Alloca, Tracker &Tracker);
-  void revert() final;
+  GenericSetter(InstrT *I, Tracker &Tracker)
+      : IRChangeBase(Tracker), I(I), OrigVal((I->*GetterFn)()) {}
+  void revert() final { (I->*SetterFn)(OrigVal); }
   void accept() final {}
 #ifndef NDEBUG
   void dump(raw_ostream &OS) const final {
     dumpCommon(OS);
-    OS << "AllocaSetAlignment";
+    OS << "GenericSetter";
   }
-  LLVM_DUMP_METHOD void dump() const final;
-#endif
-};
-
-class AllocaSetUsedWithInAlloca final : public IRChangeBase {
-  AllocaInst *Alloca;
-  bool Orig;
-
-public:
-  AllocaSetUsedWithInAlloca(AllocaInst *Alloca, Tracker &Tracker);
-  void revert() final;
-  void accept() final {}
-#ifndef NDEBUG
-  void dump(raw_ostream &OS) const final {
-    dumpCommon(OS);
-    OS << "AllocaSetUsedWithInAlloca";
+  LLVM_DUMP_METHOD void dump() const final {
+    dump(dbgs());
+    dbgs() << "\n";
   }
-  LLVM_DUMP_METHOD void dump() const final;
 #endif
 };
 
@@ -325,25 +299,6 @@ class CallBrInstSetIndirectDest : public IRChangeBase {
 #endif
 };
 
-class SetVolatile : public IRChangeBase {
-  /// This holds the properties of whether LoadInst or StoreInst was volatile
-  bool WasVolatile;
-  /// This could either be StoreInst or LoadInst
-  PointerUnion<StoreInst *, LoadInst *> StoreOrLoad;
-
-public:
-  SetVolatile(Instruction *I, Tracker &Tracker);
-  void revert() final;
-  void accept() final {}
-#ifndef NDEBUG
-  void dump(raw_ostream &OS) const final {
-    dumpCommon(OS);
-    OS << "SetVolatile";
-  }
-  LLVM_DUMP_METHOD void dump() const final;
-#endif
-};
-
 class MoveInstr : public IRChangeBase {
   /// The instruction that moved.
   Instruction *MovedI;
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -626,8 +626,11 @@ void BranchInst::dump() const {
 
 void LoadInst::setVolatile(bool V) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.isTracking())
-    Tracker.track(std::make_unique<SetVolatile>(this, Tracker));
+  if (Tracker.isTracking()) {
+    Tracker.track(std::make_unique<
+                  GenericSetter<&LoadInst::isVolatile, &LoadInst::setVolatile>>(
+        this, Tracker));
+  }
   cast<llvm::LoadInst>(Val)->setVolatile(V);
 }
 
@@ -688,8 +691,12 @@ void LoadInst::dump() const {
 
 void StoreInst::setVolatile(bool V) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.isTracking())
-    Tracker.track(std::make_unique<SetVolatile>(this, Tracker));
+  if (Tracker.isTracking()) {
+    Tracker.track(
+        std::make_unique<
+            GenericSetter<&StoreInst::isVolatile, &StoreInst::setVolatile>>(
+            this, Tracker));
+  }
   cast<llvm::StoreInst>(Val)->setVolatile(V);
 }
 
@@ -1042,8 +1049,11 @@ llvm::SmallVector<BasicBlock *, 16> CallBrInst::getIndirectDests() const {
 }
 void CallBrInst::setDefaultDest(BasicBlock *BB) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.isTracking())
-    Tracker.track(std::make_unique<CallBrInstSetDefaultDest>(this, Tracker));
+  if (Tracker.isTracking()) {
+    Tracker.track(std::make_unique<GenericSetter<&CallBrInst::getDefaultDest,
+                                                 &CallBrInst::setDefaultDest>>(
+        this, Tracker));
+  }
   cast<llvm::CallBrInst>(Val)->setDefaultDest(cast<llvm::BasicBlock>(BB->Val));
 }
 void CallBrInst::setIndirectDest(unsigned Idx, BasicBlock *BB) {
@@ -1295,22 +1305,34 @@ AllocaInst *AllocaInst::create(Type *Ty, unsigned AddrSpace,
 
 void AllocaInst::setAllocatedType(Type *Ty) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.isTracking())
-    Tracker.track(std::make_unique<AllocaSetAllocatedType>(this, Tracker));
+  if (Tracker.isTracking()) {
+    Tracker.track(
+        std::make_unique<GenericSetter<&AllocaInst::getAllocatedType,
+                                       &AllocaInst::setAllocatedType>>(
+            this, Tracker));
+  }
   cast<llvm::AllocaInst>(Val)->setAllocatedType(Ty);
 }
 
 void AllocaInst::setAlignment(Align Align) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.isTracking())
-    Tracker.track(std::make_unique<AllocaSetAlignment>(this, Tracker));
+  if (Tracker.isTracking()) {
+    Tracker.track(
+        std::make_unique<
+            GenericSetter<&AllocaInst::getAlign, &AllocaInst::setAlignment>>(
+            this, Tracker));
+  }
   cast<llvm::AllocaInst>(Val)->setAlignment(Align);
 }
 
 void AllocaInst::setUsedWithInAlloca(bool V) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.isTracking())
-    Tracker.track(std::make_unique<AllocaSetUsedWithInAlloca>(this, Tracker));
+  if (Tracker.isTracking()) {
+    Tracker.track(
+        std::make_unique<GenericSetter<&AllocaInst::isUsedWithInAlloca,
+                                       &AllocaInst::setUsedWithInAlloca>>(
+            this, Tracker));
+  }
   cast<llvm::AllocaInst>(Val)->setUsedWithInAlloca(V);
 }
 
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
@@ -204,61 +204,6 @@ void RemoveFromParent::dump() const {
 }
 #endif
 
-AllocaSetAllocatedType::AllocaSetAllocatedType(AllocaInst *Alloca,
-                                               Tracker &Tracker)
-    : IRChangeBase(Tracker), Alloca(Alloca),
-      OrigType(Alloca->getAllocatedType()) {}
-
-void AllocaSetAllocatedType::revert() { Alloca->setAllocatedType(OrigType); }
-
-#ifndef NDEBUG
-void AllocaSetAllocatedType::dump() const {
-  dump(dbgs());
-  dbgs() << "\n";
-}
-#endif // NDEBUG
-
-AllocaSetAlignment::AllocaSetAlignment(AllocaInst *Alloca, Tracker &Tracker)
-    : IRChangeBase(Tracker), Alloca(Alloca), OrigAlign(Alloca->getAlign()) {}
-
-void AllocaSetAlignment::revert() { Alloca->setAlignment(OrigAlign); }
-
-#ifndef NDEBUG
-void AllocaSetAlignment::dump() const {
-  dump(dbgs());
-  dbgs() << "\n";
-}
-#endif // NDEBUG
-
-AllocaSetUsedWithInAlloca::AllocaSetUsedWithInAlloca(AllocaInst *Alloca,
-                                                     Tracker &Tracker)
-    : IRChangeBase(Tracker), Alloca(Alloca),
-      Orig(Alloca->isUsedWithInAlloca()) {}
-
-void AllocaSetUsedWithInAlloca::revert() { Alloca->setUsedWithInAlloca(Orig); }
-
-#ifndef NDEBUG
-void AllocaSetUsedWithInAlloca::dump() const {
-  dump(dbgs());
-  dbgs() << "\n";
-}
-#endif // NDEBUG
-
-CallBrInstSetDefaultDest::CallBrInstSetDefaultDest(CallBrInst *CallBr,
-                                                   Tracker &Tracker)
-    : IRChangeBase(Tracker), CallBr(CallBr) {
-  OrigDefaultDest = CallBr->getDefaultDest();
-}
-void CallBrInstSetDefaultDest::revert() {
-  CallBr->setDefaultDest(OrigDefaultDest);
-}
-#ifndef NDEBUG
-void CallBrInstSetDefaultDest::dump() const {
-  dump(dbgs());
-  dbgs() << "\n";
-}
-#endif
-
 CallBrInstSetIndirectDest::CallBrInstSetIndirectDest(CallBrInst *CallBr,
                                                      unsigned Idx,
                                                      Tracker &Tracker)
@@ -275,35 +220,6 @@ void CallBrInstSetIndirectDest::dump() const {
 }
 #endif
 
-SetVolatile::SetVolatile(Instruction *I, Tracker &Tracker)
-    : IRChangeBase(Tracker) {
-  if (auto *Load = dyn_cast<LoadInst>(I)) {
-    WasVolatile = Load->isVolatile();
-    StoreOrLoad = Load;
-  } else if (auto *Store = dyn_cast<StoreInst>(I)) {
-    WasVolatile = Store->isVolatile();
-    StoreOrLoad = Store;
-  } else {
-    llvm_unreachable("Expected LoadInst or StoreInst");
-  }
-}
-
-void SetVolatile::revert() {
-  if (auto *Load = StoreOrLoad.dyn_cast<LoadInst *>()) {
-    Load->setVolatile(WasVolatile);
-  } else if (auto *Store = StoreOrLoad.dyn_cast<StoreInst *>()) {
-    Store->setVolatile(WasVolatile);
-  } else {
-    llvm_unreachable("Expected LoadInst or StoreInst");
-  }
-}
-#ifndef NDEBUG
-void SetVolatile::dump() const {
-  dump(dbgs());
-  dbgs() << "\n";
-}
-#endif
-
 MoveInstr::MoveInstr(Instruction *MovedI, Tracker &Tracker)
     : IRChangeBase(Tracker), MovedI(MovedI) {
   if (auto *NextI = MovedI->getNextNode())
