[llvm] r275563 - Re-submit r272891 "Prevent dangling pointer problems in BranchProbabilityInfo"

[Igor Laevsky via llvm-commits]

Author: igor.laevsky
Date: Fri Jul 15 09:31:16 2016
New Revision: 275563

URL: http://llvm.org/viewvc/llvm-project?rev=275563&view=rev
Log:
Re-submit r272891 "Prevent dangling pointer problems in BranchProbabilityInfo"

Most possibly problem was caused by the same reason as PR28400. This change
bypasses it by using CallbackVH instead of AssertingVH.

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



Modified:
    llvm/trunk/include/llvm/Analysis/BranchProbabilityInfo.h
    llvm/trunk/lib/Analysis/BranchProbabilityInfo.cpp

Modified: llvm/trunk/include/llvm/Analysis/BranchProbabilityInfo.h
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/include/llvm/Analysis/BranchProbabilityInfo.h?rev=275563&r1=275562&r2=275563&view=diff
==============================================================================
--- llvm/trunk/include/llvm/Analysis/BranchProbabilityInfo.h (original)
+++ llvm/trunk/include/llvm/Analysis/BranchProbabilityInfo.h Fri Jul 15 09:31:16 2016
@@ -15,9 +15,11 @@
 #define LLVM_ANALYSIS_BRANCHPROBABILITYINFO_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ValueHandle.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/BranchProbability.h"
@@ -116,10 +118,29 @@ public:
 
   void calculate(const Function &F, const LoopInfo &LI);
 
+  /// Forget analysis results for the given basic block.
+  void eraseBlock(const BasicBlock *BB);
+
 private:
   void operator=(const BranchProbabilityInfo &) = delete;
   BranchProbabilityInfo(const BranchProbabilityInfo &) = delete;
 
+  // We need to store CallbackVH's in order to correctly handle basic block
+  // removal.
+  class BasicBlockCallbackVH final : public CallbackVH {
+    BranchProbabilityInfo *BPI;
+    void deleted() override {
+      assert(BPI != nullptr);
+      BPI->eraseBlock(cast<BasicBlock>(getValPtr()));
+      BPI->Handles.erase(*this);
+    }
+
+  public:
+    BasicBlockCallbackVH(const Value *V, BranchProbabilityInfo *BPI=nullptr)
+        : CallbackVH(const_cast<Value *>(V)), BPI(BPI) {}
+  };
+  DenseSet<BasicBlockCallbackVH, DenseMapInfo<Value*>> Handles;
+
   // Since we allow duplicate edges from one basic block to another, we use
   // a pair (PredBlock and an index in the successors) to specify an edge.
   typedef std::pair<const BasicBlock *, unsigned> Edge;

Modified: llvm/trunk/lib/Analysis/BranchProbabilityInfo.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/lib/Analysis/BranchProbabilityInfo.cpp?rev=275563&r1=275562&r2=275563&view=diff
==============================================================================
--- llvm/trunk/lib/Analysis/BranchProbabilityInfo.cpp (original)
+++ llvm/trunk/lib/Analysis/BranchProbabilityInfo.cpp Fri Jul 15 09:31:16 2016
@@ -624,6 +624,7 @@ void BranchProbabilityInfo::setEdgeProba
                                                unsigned IndexInSuccessors,
                                                BranchProbability Prob) {
   Probs[std::make_pair(Src, IndexInSuccessors)] = Prob;
+  Handles.insert(BasicBlockCallbackVH(Src, this));
   DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << IndexInSuccessors
                << " successor probability to " << Prob << "\n");
 }
@@ -641,6 +642,14 @@ BranchProbabilityInfo::printEdgeProbabil
   return OS;
 }
 
+void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
+  for (auto I = Probs.begin(), E = Probs.end(); I != E; ++I) {
+    auto Key = I->first;
+    if (Key.first == BB)
+      Probs.erase(Key);
+  }
+}
+
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI) {
   DEBUG(dbgs() << "---- Branch Probability Info : " << F.getName()
                << " ----\n\n");