[JumpThreading] Prevent dangling pointer problems in BranchProbabilityInfo

We should update results of the BranchProbabilityInfo after removing block in JumpThreading. Otherwise 
we will get dangling pointer inside BranchProbabilityInfo cache.

Differential Revision: http://reviews.llvm.org/D20957

llvm-svn: 272891

2 files changed, 17 insertions, 5 deletions

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index 1d8d2a5e469..db083498c3a 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -182,7 +182,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
   // back edges. This works for normal cases but not for unreachable blocks as
   // they may have cycle with no back edge.
   bool EverChanged = false;
-  EverChanged |= removeUnreachableBlocks(F, LVI);
+  EverChanged |= removeUnreachableBlocks(F, LVI, BPI.get());
 
   FindLoopHeaders(F);
 
@@ -204,7 +204,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
         DEBUG(dbgs() << "  JT: Deleting dead block '" << BB->getName()
               << "' with terminator: " << *BB->getTerminator() << '\n');
         LoopHeaders.erase(BB);
-        LVI->eraseBlock(BB);
+        dropBlockAnalysisResults(BB);
         DeleteDeadBlock(BB);
         Changed = true;
         continue;
@@ -232,7 +232,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
         // for a block even if it doesn't get erased.  This isn't totally
         // awesome, but it allows us to use AssertingVH to prevent nasty
         // dangling pointer issues within LazyValueInfo.
-        LVI->eraseBlock(BB);
+        dropBlockAnalysisResults(BB);
         if (TryToSimplifyUncondBranchFromEmptyBlock(BB)) {
           Changed = true;
           // If we deleted BB and BB was the header of a loop, then the
@@ -715,7 +715,7 @@ bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
       if (LoopHeaders.erase(SinglePred))
         LoopHeaders.insert(BB);
 
-      LVI->eraseBlock(SinglePred);
+      dropBlockAnalysisResults(SinglePred);
       MergeBasicBlockIntoOnlyPred(BB);
 
       return true;
@@ -1949,3 +1949,11 @@ bool JumpThreadingPass::TryToUnfoldSelectInCurrBB(BasicBlock *BB) {
   
   return false;
 }
+
+/// dropBlockAnalysisResults - Inform relevant analyzes that BB is going to
+/// be removed. This is important in order to prevent dangling pointer problems.
+void JumpThreadingPass::dropBlockAnalysisResults(BasicBlock *BB) {
+  LVI->eraseBlock(BB);
+  if (BPI)
+    BPI->eraseBlock(BB);
+}
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index edde166b1ed..13f983231bd 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -20,6 +20,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/EHPersonalities.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
@@ -1486,7 +1487,8 @@ void llvm::removeUnwindEdge(BasicBlock *BB) {
 /// removeUnreachableBlocksFromFn - Remove blocks that are not reachable, even
 /// if they are in a dead cycle.  Return true if a change was made, false
 /// otherwise.
-bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI) {
+bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI, 
+                                   BranchProbabilityInfo *BPI) {
   SmallPtrSet<BasicBlock*, 16> Reachable;
   bool Changed = markAliveBlocks(F, Reachable);
 
@@ -1509,6 +1511,8 @@ bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI) {
         (*SI)->removePredecessor(&*BB);
     if (LVI)
       LVI->eraseBlock(&*BB);
+    if (BPI)
+      BPI->eraseBlock(&*BB);
     BB->dropAllReferences();
   }