[SandboxIR] Add debug checker to compare IR before/after a revert (#115968)

This will help us catch mistakes in change tracking. It's only enabled
when EXPENSIVE_CHECKS are enabled.

Author: slackito

llvm/include/llvm/SandboxIR/Context.h

@@ -44,11 +44,12 @@ class Context {
 
 protected:
   LLVMContext &LLVMCtx;
-  friend class Type;        // For LLVMCtx.
-  friend class PointerType; // For LLVMCtx.
-  friend class IntegerType; // For LLVMCtx.
-  friend class StructType;  // For LLVMCtx.
-  friend class Region;      // For LLVMCtx.
+  friend class Type;              // For LLVMCtx.
+  friend class PointerType;       // For LLVMCtx.
+  friend class IntegerType;       // For LLVMCtx.
+  friend class StructType;        // For LLVMCtx.
+  friend class Region;            // For LLVMCtx.
+  friend class IRSnapshotChecker; // To snapshot LLVMModuleToModuleMap.
 
   Tracker IRTracker;
 

llvm/include/llvm/SandboxIR/Instruction.h

@@ -11,6 +11,7 @@
 
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/SandboxIR/BasicBlock.h"
 #include "llvm/SandboxIR/Constant.h"

llvm/include/llvm/SandboxIR/Tracker.h

@@ -42,13 +42,12 @@
 
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StableHashing.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
-#include "llvm/IR/Module.h"
 #include "llvm/SandboxIR/Use.h"
 #include "llvm/Support/Debug.h"
 #include <memory>
-#include <regex>
 
 namespace llvm::sandboxir {
 
@@ -64,9 +63,56 @@ class SwitchInst;
 class ConstantInt;
 class ShuffleVectorInst;
 class CmpInst;
-class Module;
 class GlobalVariable;
 
+#ifndef NDEBUG
+
+/// A class that saves hashes and textual IR snapshots of functions in a
+/// SandboxIR Context, and does hash comparison when `expectNoDiff` is called.
+/// If hashes differ, it prints textual IR for both old and new versions to
+/// aid debugging.
+///
+/// This is used as an additional debug check when reverting changes to
+/// SandboxIR, to verify the reverted state matches the initial state.
+class IRSnapshotChecker {
+  Context &Ctx;
+
+  // A snapshot of textual IR for a function, with a hash for quick comparison.
+  struct FunctionSnapshot {
+    llvm::stable_hash Hash;
+    std::string TextualIR;
+  };
+
+  // A snapshot for each llvm::Function found in every module in the SandboxIR
+  // Context. In practice there will always be one module, but sandbox IR
+  // save/restore ops work at the Context level, so we must take the full state
+  // into account.
+  using ContextSnapshot = DenseMap<const llvm::Function *, FunctionSnapshot>;
+
+  ContextSnapshot OrigContextSnapshot;
+
+  // Dumps to a string the textual IR for a single Function.
+  std::string dumpIR(const llvm::Function &F) const;
+
+  // Returns a snapshot of all the modules in the sandbox IR context.
+  ContextSnapshot takeSnapshot() const;
+
+  // Compares two snapshots and returns true if they differ.
+  bool diff(const ContextSnapshot &Orig, const ContextSnapshot &Curr) const;
+
+public:
+  IRSnapshotChecker(Context &Ctx) : Ctx(Ctx) {}
+
+  /// Saves a snapshot of the current state. If there was any previous snapshot,
+  /// it will be replaced with the new one.
+  void save();
+
+  /// Checks current state against saved state, crashes if different.
+  void expectNoDiff();
+};
+
+#endif // NDEBUG
+
 /// The base class for IR Change classes.
 class IRChangeBase {
 protected:
@@ -405,14 +451,26 @@ class Tracker {
   TrackerState State = TrackerState::Disabled;
   Context &Ctx;
 
+#ifndef NDEBUG
+  IRSnapshotChecker SnapshotChecker;
+#endif
+
 public:
 #ifndef NDEBUG
   /// Helps catch bugs where we are creating new change objects while in the
   /// middle of creating other change objects.
   bool InMiddleOfCreatingChange = false;
 #endif // NDEBUG
 
-  explicit Tracker(Context &Ctx) : Ctx(Ctx) {}
+  explicit Tracker(Context &Ctx)
+      : Ctx(Ctx)
+#ifndef NDEBUG
+        ,
+        SnapshotChecker(Ctx)
+#endif
+  {
+  }
+
   ~Tracker();
   Context &getContext() const { return Ctx; }
   /// Record \p Change and take ownership. This is the main function used to

llvm/lib/SandboxIR/Tracker.cpp

@@ -10,12 +10,75 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/StructuralHash.h"
 #include "llvm/SandboxIR/Instruction.h"
 #include <sstream>
 
 using namespace llvm::sandboxir;
 
 #ifndef NDEBUG
+
+std::string IRSnapshotChecker::dumpIR(const llvm::Function &F) const {
+  std::string Result;
+  raw_string_ostream SS(Result);
+  F.print(SS, /*AssemblyAnnotationWriter=*/nullptr);
+  return Result;
+}
+
+IRSnapshotChecker::ContextSnapshot IRSnapshotChecker::takeSnapshot() const {
+  ContextSnapshot Result;
+  for (const auto &Entry : Ctx.LLVMModuleToModuleMap)
+    for (const auto &F : *Entry.first) {
+      FunctionSnapshot Snapshot;
+      Snapshot.Hash = StructuralHash(F, /*DetailedHash=*/true);
+      Snapshot.TextualIR = dumpIR(F);
+      Result[&F] = Snapshot;
+    }
+  return Result;
+}
+
+bool IRSnapshotChecker::diff(const ContextSnapshot &Orig,
+                             const ContextSnapshot &Curr) const {
+  bool DifferenceFound = false;
+  for (const auto &[F, OrigFS] : Orig) {
+    auto CurrFSIt = Curr.find(F);
+    if (CurrFSIt == Curr.end()) {
+      DifferenceFound = true;
+      dbgs() << "Function " << F->getName() << " not found in current IR.\n";
+      dbgs() << OrigFS.TextualIR << "\n";
+      continue;
+    }
+    const FunctionSnapshot &CurrFS = CurrFSIt->second;
+    if (OrigFS.Hash != CurrFS.Hash) {
+      DifferenceFound = true;
+      dbgs() << "Found IR difference in Function " << F->getName() << "\n";
+      dbgs() << "Original:\n" << OrigFS.TextualIR << "\n";
+      dbgs() << "Current:\n" << CurrFS.TextualIR << "\n";
+    }
+  }
+  // Check that Curr doesn't contain any new functions.
+  for (const auto &[F, CurrFS] : Curr) {
+    if (!Orig.contains(F)) {
+      DifferenceFound = true;
+      dbgs() << "Function " << F->getName()
+             << " found in current IR but not in original snapshot.\n";
+      dbgs() << CurrFS.TextualIR << "\n";
+    }
+  }
+  return DifferenceFound;
+}
+
+void IRSnapshotChecker::save() { OrigContextSnapshot = takeSnapshot(); }
+
+void IRSnapshotChecker::expectNoDiff() {
+  ContextSnapshot CurrContextSnapshot = takeSnapshot();
+  if (diff(OrigContextSnapshot, CurrContextSnapshot)) {
+    llvm_unreachable(
+        "Original and current IR differ! Probably a checkpointing bug.");
+  }
+}
+
 void UseSet::dump() const {
   dump(dbgs());
   dbgs() << "\n";
@@ -275,14 +338,22 @@ void CmpSwapOperands::dump() const {
 }
 #endif
 
-void Tracker::save() { State = TrackerState::Record; }
+void Tracker::save() {
+  State = TrackerState::Record;
+#if !defined(NDEBUG) && defined(EXPENSIVE_CHECKS)
+  SnapshotChecker.save();
+#endif
+}
 
 void Tracker::revert() {
   assert(State == TrackerState::Record && "Forgot to save()!");
   State = TrackerState::Disabled;
   for (auto &Change : reverse(Changes))
     Change->revert(*this);
   Changes.clear();
+#if !defined(NDEBUG) && defined(EXPENSIVE_CHECKS)
+  SnapshotChecker.expectNoDiff();
+#endif
 }
 
 void Tracker::accept() {

llvm/unittests/SandboxIR/TrackerTest.cpp

@@ -1844,3 +1844,66 @@ define void @foo(i32 %arg, float %farg) {
   Ctx.revert();
   EXPECT_FALSE(FAdd->getFastMathFlags() != OrigFMF);
 }
+
+TEST_F(TrackerTest, IRSnapshotCheckerNoChanges) {
+  parseIR(C, R"IR(
+define i32 @foo(i32 %arg) {
+  %add0 = add i32 %arg, %arg
+  ret i32 %add0
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+
+  [[maybe_unused]] auto *F = Ctx.createFunction(&LLVMF);
+  sandboxir::IRSnapshotChecker Checker(Ctx);
+  Checker.save();
+  Checker.expectNoDiff();
+}
+
+TEST_F(TrackerTest, IRSnapshotCheckerDiesWithUnexpectedChanges) {
+  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::IRSnapshotChecker Checker(Ctx);
+  Checker.save();
+  Add1->setOperand(1, Add0);
+  EXPECT_DEATH(Checker.expectNoDiff(), "Found IR difference");
+}
+
+TEST_F(TrackerTest, IRSnapshotCheckerSaveMultipleTimes) {
+  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::IRSnapshotChecker Checker(Ctx);
+  Checker.save();
+  Add1->setOperand(1, Add0);
+  // Now IR differs from the last snapshot. Let's take a new snapshot.
+  Checker.save();
+  // The new snapshot should have replaced the old one, so this should succeed.
+  Checker.expectNoDiff();
+}