1 files changed, 11 insertions, 3 deletions
diff --git a/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp b/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp
index 7d1e630e6e80..9fee57c54b85 100644
--- a/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp
+++ b/llvm/lib/Analysis/InstructionPrecedenceTracking.cpp
@@ -19,11 +19,15 @@
 
 #include "llvm/Analysis/InstructionPrecedenceTracking.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
 
+#define DEBUG_TYPE "ipt"
+STATISTIC(NumInstScanned, "Number of insts scanned while updating ibt");
+
 #ifndef NDEBUG
 static cl::opt<bool> ExpensiveAsserts(
     "ipt-expensive-asserts",
@@ -64,11 +68,13 @@ bool InstructionPrecedenceTracking::isPreceededBySpecialInstruction(
 
 void InstructionPrecedenceTracking::fill(const BasicBlock *BB) {
   FirstSpecialInsts.erase(BB);
-  for (auto &I : *BB)
+  for (auto &I : *BB) {
+    NumInstScanned++;
     if (isSpecialInstruction(&I)) {
       FirstSpecialInsts[BB] = &I;
       return;
     }
+  }
 
   // Mark this block as having no special instructions.
   FirstSpecialInsts[BB] = nullptr;
@@ -107,8 +113,10 @@ void InstructionPrecedenceTracking::insertInstructionTo(const Instruction *Inst,
 }
 
 void InstructionPrecedenceTracking::removeInstruction(const Instruction *Inst) {
-  if (isSpecialInstruction(Inst))
-    FirstSpecialInsts.erase(Inst->getParent());
+  auto *BB = Inst->getParent();
+  assert(BB && "must be called before instruction is actually removed");
+  if (FirstSpecialInsts.count(BB) && FirstSpecialInsts[BB] == Inst)
+    FirstSpecialInsts.erase(BB);
 }
 
 void InstructionPrecedenceTracking::removeUsersOf(const Instruction *Inst) {