[Chkpnt] Adds IR Checker that compares the saved and restored states.

[vporpo]

Summary

This patch adds a simple IR checker that compares the IR dump of the tracked
components when tracking initially took place and after restore() was called.
This should run only in debug builds, and should help us catch errors in the
checkpointing infrastructure.

llvm/include/llvm/IR/Checkpoint.h

 namespace llvm {
 
 class raw_ostream;
 class BasicBlock;
 class CheckpointTracker;
 
 class Checkpoint {
   CheckpointTracker &ChkpntTracker;
+  bool RunVerifier;
   /// No copies allowed because going out of scope checks if we restored or
   /// accepted the changes.
   Checkpoint(Checkpoint &) = delete;
   void operator=(const Checkpoint &) = delete;
 
 public:
-  Checkpoint(CheckpointTracker &ChkpntTracker) : ChkpntTracker(ChkpntTracker) {}
+  /// If \p RunVerifier is true we run expensive checks to compare the modules
+  /// state between save() and restore(). These check if checkpointing works
+  /// correctly by comparing the saved and restored states. The checks only run
+  /// in a debug build, but if the saved components are large they can be
+  /// expensive.
+  Checkpoint(CheckpointTracker &ChkpntTracker, bool RunVerifier)
+      : ChkpntTracker(ChkpntTracker), RunVerifier(RunVerifier) {}
   ~Checkpoint();
 
   /// \p MaxNumOfTrackedChanges is used for debugging to help diagnose cases
   /// were the user forgets to accept() or restore(). It will cause a crash
   /// if we record more changes that this number.
   void setMaxNumOfTrackedChanges(uint32_t MaxNumOfTrackedChanges);
   /// Activates checkpointing and starts tracking changes made to \p BB. This
   /// includes all instructions in the block and their state, including their

llvm/include/llvm/IR/CheckpointTracker.h

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/CheckpointCommon.h"
 #include "llvm/Support/Debug.h"
 #include <memory>
 #include <regex>
 
 namespace llvm {
 class BasicBlock;
+
+#ifndef NDEBUG
+/// Helper abstract class that saves the textual representation of the IR upon
+/// construction and compares against it when `expectNoDiff()` is called.
+/// Each component should implement its own custom checker class that inherits
+/// from this.
+class IRChecker {
+protected:
+  /// Holds the initial IR dump.
+  std::string OriginalIR;
+  /// Helper string that helps skip the list of predecessors. This is to avoid
+  /// false positives due to differences in the order of users.
+  std::string PredsRegexStr = "; preds = .*\n";
+  /// Holds the std::regex object created from `PredsRegexStr`.
+  std::regex PredsRegex;
+  /// Prints a simple line diff between \p OrigIR and \p CurrIR.
+  static void showDiff(const std::string &OrigIR, const std::string &CurrIR);
+
+public:
+  IRChecker() = default;
+  virtual ~IRChecker() = default;
+  IRChecker(const IRChecker &) = delete;
+  IRChecker(IRChecker &&) = delete;
+  /// This saves the IR state into `OriginalIR`.
+  virtual void save() = 0;
+  /// \Returns the dump of the original IR.
+  const std::string &origIR() const { return OriginalIR; }
+  /// \Returns the dump of the current IR.
+  virtual std::string currIR() const = 0;
+  /// Crashes if there is a difference between the original and current IR.
+  void expectNoDiff() const;
+};
+
+/// Checker for a single BasicBlock.
+class BBChecker : public IRChecker {
+  BasicBlock *BB = nullptr;
+  bool SkipPreds;
+
+  std::string dumpBB() const;
+
+public:
+  /// If SkipPreds is true we remove the "; preds = .*" string from the dumps.
+  /// This helps avoids false positives since Checkpoint cannot currently
+  /// preserve the order of users.
+  BBChecker(BasicBlock *BB, bool SkipPreds = true);
+  /// We dump the IR of the BasicBlock into `OriginalIR`.
+  void save() override;
+  /// \Returns an IR dump of the current state of the BasicBlock.
+  std::string currIR() const override;
+};
+#endif // NDEBUG
+
 class ChangeBase;
 class Instruction;
 class CheckpointTracker;
 
 /// A simple guard class that deactivates checkpointing on construction and
 /// reactivates it on destruction.
 class CheckpointGuard {
   CheckpointTracker *Chkpnt;

 /// This is the main class for the checkpointing internals. This is where
 /// the changes get recorded.
 class CheckpointTracker {
   /// The sequence of changes applied to the IR in the order they take place.
   SmallVector<std::unique_ptr<ChangeBase>, 32> Changes;
   /// The set of components currently being tracked.
   DenseSet<ChkpntComponent> ComponentsTracked;
 
+#ifndef NDEBUG
+  /// Controls whether we are running the verifier in rollback().
+  bool RunVerifier = false;
+  /// The verifier.
+  DenseMap<ChkpntComponent, std::unique_ptr<IRChecker>> IRCheckers;
+#endif // NDEBUG
+
   /// This is true while checkpointing is active.
   bool Active = false;
   friend class CheckpointGuard; // Needs access to `Active`.
 
   /// A limit to the number of changes we will record. Going over the limit
   /// causes a crash. This is useful for debugging, as it will catch cases where
   /// we collect too many changes, which would suggest that the user forgot to
   /// accept() or restore().
#endif                     // NDEBUG
   /// Called by V->takeName(FromV).
   void takeName(Value *V, Value *FromV);
   /// Gets called by ~Value() before destroyValueName().
   void destroyName(Value *V);
 
   // Main API functions. These are called by `Checkpoint`.
 
   /// Start tracking IR changes for \p Component from this point on.
-  void trackComponent(ChkpntComponent Component);
+  /// \p RunVerifier runs checks that compare component's state after
+  /// restore and before tracking started (only in DEBUG builds).
+  void trackComponent(ChkpntComponent Component, bool Runverifier);
   /// Override the maximum number of tracked changes.
   void setMaxNumOfTrackedChanges(uint32_t MaxNumOfTrackedChanges);
   /// Accept all changes.
   void acceptComponents();
   /// Reverts all changes.
   void restoreComponents();
 
 #ifndef NDEBUG

llvm/include/llvm/IR/LLVMContext.h

#undef LLVM_FIXED_MD_KIND
   /// the default value.
   void setOpaquePointers(bool Enable) const;
 
   /// Whether typed pointers are supported. If false, all pointers are opaque.
   bool supportsTypedPointers() const;
 
   /// \Returns the checkpoint handle that allows us to save/restore the state of
   /// IR components being tracked.
-  Checkpoint getCheckpoint();
+  Checkpoint getCheckpoint(bool RunVerifier = true);
 
 private:
   // Module needs access to the add/removeModule methods.
   friend class Module;
 
   /// addModule - Register a module as being instantiated in this context.  If
   /// the context is deleted, the module will be deleted as well.
   void addModule(Module*);

llvm/lib/IR/Checkpoint.cpp

Checkpoint::~Checkpoint() {
   assert(ChkpntTracker.empty() &&
          "Missing call to Checkpoint::accept() or Checkpoint::restore()");
 }
 
 void Checkpoint::setMaxNumOfTrackedChanges(uint32_t MaxNumOfTrackedChanges) {
   ChkpntTracker.setMaxNumOfTrackedChanges(MaxNumOfTrackedChanges);
 }
 
-void Checkpoint::track(BasicBlock *BB) { ChkpntTracker.trackComponent(BB); }
+void Checkpoint::track(BasicBlock *BB) {
+  ChkpntTracker.trackComponent(BB, RunVerifier);
+}
 
 void Checkpoint::restore() { ChkpntTracker.restoreComponents(); }
 
 void Checkpoint::accept() { ChkpntTracker.acceptComponents(); }
 
 bool Checkpoint::empty() const { return ChkpntTracker.empty(); }
 
 #ifndef NDEBUG
 uint32_t Checkpoint::size() const { return ChkpntTracker.size(); }
 
 void Checkpoint::dump(raw_ostream &OS) const { ChkpntTracker.dump(OS); }
 
 void Checkpoint::dump() const { dump(dbgs()); }
 #endif // NDEBUG

llvm/lib/IR/CheckpointTracker.cpp

 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include <sstream>
 
 using namespace llvm;
 using namespace std;
 
+#ifndef NDEBUG
+void IRChecker::showDiff(const std::string &OrigIR, const std::string &CurrIR) {
+  // Show the first line that differes.
+  std::stringstream OrigSS(OrigIR);
+  std::stringstream CurrSS(CurrIR);
+  std::string OrigLine;
+  std::string CurrLine;
+  SmallVector<std::string> Context;
+  static constexpr const uint32_t MaxContext = 3;
+  while (OrigSS.good() && CurrSS.good()) {
+    std::getline(OrigSS, OrigLine);
+    std::getline(CurrSS, CurrLine);
+    if (CurrLine != OrigLine) {
+      // Print context.
+      for (const std::string &ContextLine : Context)
+        dbgs() << "  " << ContextLine << "\n";
+      // Print the line difference.
+      dbgs() << "- " << OrigLine << "\n";
+      dbgs() << "+ " << CurrLine << "\n";
+    } else {
+      // Lazy way to maintain context. Performance of this code does not matter.
+      Context.push_back(OrigLine);
+      if (Context.size() > MaxContext)
+        Context.erase(Context.begin());
+    }
+  }
+  // If one file is larger than the other print line in the larger one.
+  if (!OrigSS.good() && CurrSS.good()) {
+    std::getline(CurrSS, CurrLine);
+    dbgs() << "+ " << CurrLine << "\n";
+  }
+  if (OrigSS.good() && !CurrSS.good()) {
+    std::getline(OrigSS, OrigLine);
+    dbgs() << "+ " << OrigLine << "\n";
+  }
+}
+
+void IRChecker::expectNoDiff() const {
+  const std::string &OrigIR = origIR();
+  std::string CurrIR = currIR();
+  bool Same = OrigIR == CurrIR;
+  if (!Same) {
+    showDiff(OrigIR, CurrIR);
+    llvm_unreachable(
+        "Original and current IR differ! Possibly a Checkpointing bug.");
+  }
+}
+
+BBChecker::BBChecker(BasicBlock *BB, bool SkipPreds) : BB(BB) {}
+
+std::string BBChecker::dumpBB() const {
+  std::string TmpStr;
+  raw_string_ostream SS(TmpStr);
+  BB->print(SS, /*AssemblyAnnotationWriter=*/nullptr);
+  return SkipPreds
+             ? std::regex_replace(TmpStr, PredsRegex, "; preds = <removed>\n")
+             : TmpStr;
+}
+
+void BBChecker::save() {
+  assert(BB != nullptr && "BB not set!");
+  if (SkipPreds)
+    PredsRegex = std::regex(PredsRegexStr);
+  OriginalIR = dumpBB();
+}
+
+std::string BBChecker::currIR() const { return dumpBB(); }
+#endif // NDEBUG
+
 CheckpointGuard::CheckpointGuard(bool NewState, CheckpointTracker *Chkpnt)
     : Chkpnt(Chkpnt), LastState(Chkpnt->isActive()) {
   Chkpnt->Active = NewState;
 }
 
 CheckpointGuard::~CheckpointGuard() { Chkpnt->Active = LastState; }
 
 void CheckpointTracker::clear() {
if (!Parent)
     return;
   // If changing name of a BB, we consider its parent to be the BB itself.
   if (isa<BasicBlock>(V) && trackingComponent(cast<BasicBlock>(V)))
     Parent = cast<BasicBlock>(V);
   if (trackingComponent(*Parent))
     Changes.push_back(make_unique<DestroyName>(V, this));
 }
 
-void CheckpointTracker::trackComponent(ChkpntComponent Component) {
+void CheckpointTracker::trackComponent(ChkpntComponent Component,
+                                       bool RunVerifier) {
+  this->RunVerifier = RunVerifier;
   assert(!trackingComponent(Component) && "Already tracking component");
   ComponentsTracked.insert(Component);
+#ifndef NDEBUG
+  if (RunVerifier) {
+    if (BasicBlock **BBPtr = std::get_if<BasicBlock *>(&Component))
+      IRCheckers[Component] = make_unique<BBChecker>(*BBPtr);
+    else
+      llvm_unreachable("Unimplemented IRChecker for `Component`");
+    IRCheckers[Component]->save();
+  }
+#endif // NDEBUG
   Active = true;
 }
 
 void CheckpointTracker::setMaxNumOfTrackedChanges(
     uint32_t MaxNumOfTrackedChanges) {
   MaxNumChanges = MaxNumOfTrackedChanges;
 }
 
 void CheckpointTracker::acceptComponents() {
   {
     // Deactivate tracking temporarily to make sure that any functions used by
     // `apply()` are not being tracked.
     auto DisableTrackingGuard = disable();
     for (auto &ChangePtr : Changes)
       ChangePtr->apply();
   }
   Changes.clear();
   ComponentsTracked.clear();
   Active = false;
+#ifndef NDEBUG
+  IRCheckers.clear();
+#endif // NDEBUG
 }
 
 void CheckpointTracker::restoreComponents() {
   assert(Active && "Trying to restore() without having started tracking");
   {
     // Deactivate tracking temporarily to make sure that any functions used by
     // `revert()` are not being tracked.
     auto DisableTrackingGuard = disable();
     // Iterate through the changes in reverse and revert if they match
     // Component.
     for (auto &ChangePtr : reverse(Changes))
       ChangePtr->revert();
   }
   Changes.clear();
   ComponentsTracked.clear();
   Active = false;
+#ifndef NDEBUG
+  if (RunVerifier) {
+    for (ChkpntComponent &Component : ComponentsTracked) {
+      auto &CheckerPtr = IRCheckers[Component];
+      CheckerPtr->expectNoDiff();
+      IRCheckers.erase(Component);
+      break;
+    }
+  }
+#endif // NDEBUG
 }
 
 CheckpointTracker::CheckpointTracker() {}
 
 CheckpointTracker::~CheckpointTracker() {}
 
 bool CheckpointTracker::trackingComponent(
     const ChkpntComponent &Component) const {

llvm/lib/IR/LLVMContext.cpp

 void LLVMContext::setOpaquePointers(bool Enable) const {
   pImpl->setOpaquePointers(Enable);
 }
 
 bool LLVMContext::supportsTypedPointers() const {
   return !pImpl->getOpaquePointers();
 }
 
-Checkpoint LLVMContext::getCheckpoint() { return Checkpoint(ChkpntTracker); }
+Checkpoint LLVMContext::getCheckpoint(bool RunVerifier) {
+  return Checkpoint(ChkpntTracker, RunVerifier);
+}

llvm/unittests/IR/CheckpointTest.cpp

 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = &*F->begin();
   auto *Instr = &*std::next(BB0->begin(), 0);
 
 #ifndef NDEBUG
   EXPECT_DEATH(
       {
-        Checkpoint Chkpnt = C.getCheckpoint();
+        Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
         Chkpnt.track(BB0);
         Instr->setName("new");
       },
       ".*");
 #endif
 }
 
 TEST(CheckpointTest, MaxNumOfTrackedChanges) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"(
 define void @foo(i32 %a, i32 %b) {
 bb0:
   %instr = add i32 %a, %b
   ret void
 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = &*F->begin();
   auto *Instr = &*std::next(BB0->begin(), 0);
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0);
   Chkpnt.setMaxNumOfTrackedChanges(1);
   Instr->setName("change1");
 #ifndef NDEBUG
   // This should crash as we exceeded the maximum number of tracked changes.
   EXPECT_DEATH({ Instr->setName("change2"); }, ".*");
 #endif
   Chkpnt.accept();
 }
 
 TEST(CheckpointTest, BB_SetNameInstr) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"(
 define void @foo(i32 %a, i32 %b) {
 bb0:
   %instr = add i32 %a, %b
   ret void
 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = getBBWithName(F, "bb0");
   auto *Instr = &*std::next(BB0->begin(), 0);
 
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0);
   Instr->setName("new");
   EXPECT_NE(Instr->getName(), "instr");
   EXPECT_FALSE(Chkpnt.empty());
   Chkpnt.restore();
 
   EXPECT_EQ(Instr->getName(), "instr");
 }

 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = getBBWithName(F, "bb0");
   BasicBlock *BB1 = getBBWithName(F, "bb1");
   auto *Instr1 = &*std::next(BB0->begin(), 0);
   auto *Instr2 = &*std::next(BB1->begin(), 0);
 
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0, BB1);
   Instr1->setName("new1");
   Instr2->setName("new2");
   EXPECT_NE(Instr1->getName(), "instr1");
   EXPECT_NE(Instr2->getName(), "instr2");
   EXPECT_FALSE(Chkpnt.empty());
   Chkpnt.restore();
   EXPECT_EQ(Instr1->getName(), "instr1");

 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = getBBWithName(F, "bb0");
   BasicBlock *BB1 = getBBWithName(F, "bb1");
   auto *Instr1 = &*std::next(BB0->begin(), 0);
   auto *Instr2 = &*std::next(BB1->begin(), 0);
 
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0, BB1);
   Instr1->setName("new1");
   Instr2->setName("new2");
   EXPECT_NE(Instr1->getName(), "instr1");
   EXPECT_NE(Instr2->getName(), "instr2");
   EXPECT_FALSE(Chkpnt.empty());
   Chkpnt.restore();
   EXPECT_EQ(Instr1->getName(), "instr1");
bb0:
   %instr = add i32 %a, %b
   ret void
 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = getBBWithName(F, "bb0");
   auto *Instr = &*std::next(BB0->begin(), 0);
 
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0);
   Instr->setName("new");
   EXPECT_FALSE(Chkpnt.empty());
   Chkpnt.accept();
   EXPECT_EQ(Instr->getName(), "new");
 }
 
 TEST(CheckpointTest, BB_TakeNameInstr) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"(
 define void @foo(i32 %a, i32 %b) {
 bb0:
   %instr1 = add i32 %a, %b
   %instr2 = add i32 %a, %b
   ret void
 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = &*F->begin();
   auto It = BB0->begin();
   auto *Instr1 = &*It++;
   auto *Instr2 = &*It++;
 
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0);
   Instr1->takeName(Instr2);
   EXPECT_EQ(Instr1->getName(), "instr2");
   EXPECT_FALSE(Instr2->hasName());
   EXPECT_FALSE(Chkpnt.empty());
   Chkpnt.restore();
 
   EXPECT_EQ(Instr1->getName(), "instr1");
   EXPECT_EQ(Instr2->getName(), "instr2");
 }
 
 TEST(CheckpointTest, BB_SetNameBB) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"(
 define void @foo(i32 %a, i32 %b) {
 bb0:
   %instr = add i32 %a, %b
   ret void
 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = &*F->begin();
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0);
   BB0->setName("NEWNAME");
   EXPECT_NE(BB0->getName(), "bb0");
   EXPECT_FALSE(Chkpnt.empty());
   Chkpnt.restore();
   EXPECT_EQ(BB0->getName(), "bb0");
 }
 
 TEST(CheckpointTest, BB_SetNameFn) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"(
 define void @foo(i32 %a, i32 %b) {
 bb0:
   %instr = add i32 %a, %b
   ret void
 }
 )");
   Function *F = &*M->begin();
   BasicBlock *BB0 = getBBWithName(F, "bb0");
 
-  Checkpoint Chkpnt = C.getCheckpoint();
+  Checkpoint Chkpnt = C.getCheckpoint(/*RunVerifier=*/true);
   Chkpnt.track(BB0);
   F->setName("bar");
   EXPECT_NE(F->getName(), "foo");
   // We don't save the state of F because we are only tracking BB0
   EXPECT_TRUE(Chkpnt.empty());
   Chkpnt.restore();
   // We did not save the state of F so rollback should not touch F's name.
   EXPECT_EQ(F->getName(), "bar");
 }
 
 } // namespace
 } // namespace llvm