2 files changed, 48 insertions, 0 deletions

diff --git a/llvm/include/llvm/Analysis/InstructionPrecedenceTracking.h b/llvm/include/llvm/Analysis/InstructionPrecedenceTracking.h
index 0a82eece50f..8c70f64a366 100644
--- a/llvm/include/llvm/Analysis/InstructionPrecedenceTracking.h
+++ b/llvm/include/llvm/Analysis/InstructionPrecedenceTracking.h
@@ -37,6 +37,15 @@ class InstructionPrecedenceTracking {
   // Fills information about the given block's special instructions.
   void fill(const BasicBlock *BB);
 
+  #ifndef NDEBUG
+  /// Asserts whether or not the contents of this tracking is up-to-date. It can
+  /// be used to detect situations where we failed to invalidate the map
+  /// properly. The behavior of request to an invalid tracking is undefined, and
+  /// we should avoid such situations. It is slow and should only be called in
+  /// debug mode.
+  void validate() const;
+#endif
+
 protected:
   InstructionPrecedenceTracking(DominatorTree *DT)
       : OI(OrderedInstructions(DT)) {}
diff --git a/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp b/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp
index d0457715196..8e96809ef9f 100644
--- a/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp
+++ b/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp
@@ -25,6 +25,11 @@ using namespace llvm;
 
 const Instruction *InstructionPrecedenceTracking::getFirstSpecialInstruction(
     const BasicBlock *BB) {
+#ifndef NDEBUG
+  // Make sure that we are making this request to tracking which is up-to-date.
+  validate();
+#endif
+
   if (!KnownBlocks.count(BB))
     fill(BB);
   auto *FirstICF = FirstSpecialInsts.lookup(BB);
@@ -56,6 +61,36 @@ void InstructionPrecedenceTracking::fill(const BasicBlock *BB) {
   KnownBlocks.insert(BB);
 }
 
+#ifndef NDEBUG
+void InstructionPrecedenceTracking::validate() const {
+  unsigned NumNoSpecialBlocks = 0;
+  // Check that for every known block we have something cached for it.
+  for (auto *BB : KnownBlocks) {
+    auto It = FirstSpecialInsts.find(BB);
+    bool BlockHasSpecialInsns = false;
+    for (const Instruction &Insn : *BB) {
+      if (isSpecialInstruction(&Insn)) {
+        assert(It != FirstSpecialInsts.end() &&
+               "Blocked marked as known but we have no cached value for it!");
+        assert(It->second == &Insn &&
+               "Cached first special instruction is wrong!");
+        BlockHasSpecialInsns = true;
+        break;
+      }
+    }
+    if (!BlockHasSpecialInsns) {
+      assert(It == FirstSpecialInsts.end() &&
+             "Block is marked as having special instructions but in fact it "
+             "has none!");
+      ++NumNoSpecialBlocks;
+    }
+  }
+
+  assert(KnownBlocks.size() == NumNoSpecialBlocks + FirstSpecialInsts.size() &&
+         "We don't have info for some blocks?");
+}
+#endif
+
 void InstructionPrecedenceTracking::invalidateBlock(const BasicBlock *BB) {
   OI.invalidateBlock(BB);
   FirstSpecialInsts.erase(BB);
@@ -67,6 +102,10 @@ void InstructionPrecedenceTracking::clear() {
     OI.invalidateBlock(It.first);
   FirstSpecialInsts.clear();
   KnownBlocks.clear();
+#ifndef NDEBUG
+  // The map should be valid after clearing (at least empty).
+  validate();
+#endif
 }
 
 bool ImplicitControlFlowTracking::isSpecialInstruction(