[BranchProbabilityInfo] Remove block handles in eraseBlock()

[yrouban]

Summary

BranchProbabilityInfo::eraseBlock() is a public method and can be called without deleting the block itself.
I believe this method should remove the correspondent tracking handle from BranchProbabilityInfo::Handles along with the probabilities of the block. So I propose to move Handles.erase() call to eraseBlock().
In setEdgeProbability() we need to add the block handle only once.

[kazu]

LGTM. Thanks!

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

216

While you are at this, could you make BranchProbabilityInfo::Handles private?

[yrouban]

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

216

It is private.

[kazu]

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

216

Sorry, I missed that. Thank you for checking!

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

private:
   // 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) {}
   };
 
   /// Pair of Loop and SCC ID number. Used to unify handling of normal and

llvm/lib/Analysis/BranchProbabilityInfo.cpp

 /// Set the edge probability for all edges at once.
 void BranchProbabilityInfo::setEdgeProbability(
     const BasicBlock *Src, const SmallVectorImpl<BranchProbability> &Probs) {
   assert(Src->getTerminator()->getNumSuccessors() == Probs.size());
   eraseBlock(Src); // Erase stale data if any.
   if (Probs.size() == 0)
     return; // Nothing to set.
 
+  Handles.insert(BasicBlockCallbackVH(Src, this));
   uint64_t TotalNumerator = 0;
   for (unsigned SuccIdx = 0; SuccIdx < Probs.size(); ++SuccIdx) {
     this->Probs[std::make_pair(Src, SuccIdx)] = Probs[SuccIdx];
-    Handles.insert(BasicBlockCallbackVH(Src, this));
     LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << SuccIdx
                       << " successor probability to " << Probs[SuccIdx]
                       << "\n");
     TotalNumerator += Probs[SuccIdx].getNumerator();
   }
 
   // Because of rounding errors the total probability cannot be checked to be
   // 1.0 exactly. That is TotalNumerator == BranchProbability::getDenominator.

 void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
   // Note that we cannot use successors of BB because the terminator of BB may
   // have changed when eraseBlock is called as a BasicBlockCallbackVH callback.
   // Instead we remove prob data for the block by iterating successors by their
   // indices from 0 till the last which exists. There could not be prob data for
   // a pair (BB, N) if there is no data for (BB, N-1) because the data is always
   // set for all successors from 0 to M at once by the method
   // setEdgeProbability().
+  Handles.erase(BasicBlockCallbackVH(BB, this));
   for (unsigned I = 0;; ++I) {
     auto MapI = Probs.find(std::make_pair(BB, I));
     if (MapI == Probs.end()) {
       assert(Probs.count(std::make_pair(BB, I + 1)) == 0 &&
              "Must be no more successors");
       return;
     }
     Probs.erase(MapI);