Index: llvm/include/llvm/Analysis/BranchProbabilityInfo.h
===================================================================
--- llvm/include/llvm/Analysis/BranchProbabilityInfo.h
+++ llvm/include/llvm/Analysis/BranchProbabilityInfo.h
@@ -62,8 +62,7 @@
   }
 
   BranchProbabilityInfo(BranchProbabilityInfo &&Arg)
-      : Probs(std::move(Arg.Probs)), MaxSuccIdx(std::move(Arg.MaxSuccIdx)),
-        LastF(Arg.LastF),
+      : Probs(std::move(Arg.Probs)), LastF(Arg.LastF),
         PostDominatedByUnreachable(std::move(Arg.PostDominatedByUnreachable)),
         PostDominatedByColdCall(std::move(Arg.PostDominatedByColdCall)) {}
 
@@ -73,7 +72,6 @@
   BranchProbabilityInfo &operator=(BranchProbabilityInfo &&RHS) {
     releaseMemory();
     Probs = std::move(RHS.Probs);
-    MaxSuccIdx = std::move(RHS.MaxSuccIdx);
     PostDominatedByColdCall = std::move(RHS.PostDominatedByColdCall);
     PostDominatedByUnreachable = std::move(RHS.PostDominatedByUnreachable);
     return *this;
@@ -124,16 +122,6 @@
   raw_ostream &printEdgeProbability(raw_ostream &OS, const BasicBlock *Src,
                                     const BasicBlock *Dst) const;
 
-protected:
-  /// Set the raw edge probability for the given edge.
-  ///
-  /// This allows a pass to explicitly set the edge probability for an edge. It
-  /// can be used when updating the CFG to update and preserve the branch
-  /// probability information. Read the implementation of how these edge
-  /// probabilities are calculated carefully before using!
-  void setEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors,
-                          BranchProbability Prob);
-
 public:
   /// Set the raw probabilities for all edges from the given block.
   ///
@@ -275,9 +263,6 @@
 
   DenseMap<Edge, BranchProbability> Probs;
 
-  // The maximum successor index ever entered for a given basic block.
-  DenseMap<const BasicBlock *, unsigned> MaxSuccIdx;
-
   /// Track the last function we run over for printing.
   const Function *LastF = nullptr;
 
Index: llvm/lib/Analysis/BranchProbabilityInfo.cpp
===================================================================
--- llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -1031,7 +1031,6 @@
 
 void BranchProbabilityInfo::releaseMemory() {
   Probs.clear();
-  MaxSuccIdx.clear();
   Handles.clear();
 }
 
@@ -1093,6 +1092,10 @@
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
                                           unsigned IndexInSuccessors) const {
   auto I = Probs.find(std::make_pair(Src, IndexInSuccessors));
+  assert((Probs.end() == Probs.find(std::make_pair(Src, 0))) ==
+             (Probs.end() == I) &&
+         "Probability for I-th successor must always be defined along with the "
+         "probability for the first successor");
 
   if (I != Probs.end())
     return I->second;
@@ -1127,31 +1130,21 @@
   return FoundProb ? Prob : BranchProbability(EdgeCount, succ_num);
 }
 
-/// Set the edge probability for a given edge specified by PredBlock and an
-/// index to the successors.
-void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src,
-                                               unsigned IndexInSuccessors,
-                                               BranchProbability Prob) {
-  Probs[std::make_pair(Src, IndexInSuccessors)] = Prob;
-  Handles.insert(BasicBlockCallbackVH(Src, this));
-  LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> "
-                    << IndexInSuccessors << " successor probability to " << Prob
-                    << "\n");
-
-  unsigned &SuccIdx = MaxSuccIdx[Src];
-  SuccIdx = std::max(SuccIdx, IndexInSuccessors);
-}
-
 /// 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.
 
   uint64_t TotalNumerator = 0;
   for (unsigned SuccIdx = 0; SuccIdx < Probs.size(); ++SuccIdx) {
-    setEdgeProbability(Src, SuccIdx, Probs[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();
   }
 
@@ -1179,16 +1172,20 @@
 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.
-  auto It = MaxSuccIdx.find(BB);
-  if (It == MaxSuccIdx.end())
-    return;
-
-  for (unsigned I = 0, E = It->second; I <= E; ++I) {
+  // Instead we remove prob data for the block by iterating successors by their
+  // numbers from 0 till the last which exist. 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 public method
+  // setEdgeProbability().
+  for (unsigned I = 0; true; ++I) {
     auto MapI = Probs.find(std::make_pair(BB, I));
-    if (MapI != Probs.end())
-      Probs.erase(MapI);
+    if (MapI == Probs.end()) {
+      assert(Probs.find(std::make_pair(BB, I + 1)) == Probs.end() &&
+             "Must be no more successors");
+      return;
+    }
+    Probs.erase(MapI);
   }
-  MaxSuccIdx.erase(BB);
 }
 
 void BranchProbabilityInfo::calculate(const Function &F, const LoopInfo &LI,