From 681d6c711f2743b69eaed5b3708e7c4bf86489c2 Mon Sep 17 00:00:00 2001
From: Yevgeny Rouban <yrouban@azul.com>
Date: Fri, 6 Nov 2020 14:46:34 +0700
Subject: [PATCH] [BranchProbabilityInfo] Introduce method
 copyEdgeProbabilities(). NFC

A new method is introduced to allow bulk copy of outgoing edge
probabilities from one block to another. This can be useful when
a block is cloned from another one and we do not know if there
are edge probabilities set for the original block or not.
Copying outside of the BranchProbabilityInfo class makes the user
unconditionally set the cloned block's edge probabilities even if
they are unset for the original block.

Reviewed By: MaskRay

Differential Revision: https://reviews.llvm.org/D90839
---
 .../llvm/Analysis/BranchProbabilityInfo.h     |  6 ++++++
 llvm/lib/Analysis/BranchProbabilityInfo.cpp   | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
index 78fe4adb5ebc..64c3da80f6ea 100644
--- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
+++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
@@ -131,6 +131,12 @@ public:
   void setEdgeProbability(const BasicBlock *Src,
                           const SmallVectorImpl<BranchProbability> &Probs);
 
+  /// Copy outgoing edge probabilities from \p Src to \p Dst.
+  ///
+  /// This allows to keep probabilities unset for the destination if they were
+  /// unset for source.
+  void copyEdgeProbabilities(BasicBlock *Src, BasicBlock *Dst);
+
   static BranchProbability getBranchProbStackProtector(bool IsLikely) {
     static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20);
     return IsLikely ? LikelyProb : LikelyProb.getCompl();
diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index ad8817d1b1a6..37b388cd9a8b 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -1157,6 +1157,25 @@ void BranchProbabilityInfo::setEdgeProbability(
   assert(TotalNumerator >= BranchProbability::getDenominator() - Probs.size());
 }
 
+void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src,
+                                                  BasicBlock *Dst) {
+  eraseBlock(Dst); // Erase stale data if any.
+  unsigned NumSuccessors = Src->getTerminator()->getNumSuccessors();
+  assert(NumSuccessors == Dst->getTerminator()->getNumSuccessors());
+  if (NumSuccessors == 0)
+    return; // Nothing to set.
+  if (this->Probs.find(std::make_pair(Src, 0)) == this->Probs.end())
+    return; // No probability is set for edges from Src. Keep the same for Dst.
+
+  Handles.insert(BasicBlockCallbackVH(Dst, this));
+  for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx) {
+    auto Prob = this->Probs[std::make_pair(Src, SuccIdx)];
+    this->Probs[std::make_pair(Dst, SuccIdx)] = Prob;
+    LLVM_DEBUG(dbgs() << "set edge " << Dst->getName() << " -> " << SuccIdx
+                      << " successor probability to " << Prob << "\n");
+  }
+}
+
 raw_ostream &
 BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
                                             const BasicBlock *Src,