Fix a quadratic algorithm in MachineBranchProbabilityInfo.

The getSumForBlock function was quadratic in the number of successors
because getSuccWeight would perform a linear search for an already known
iterator.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@161460 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 7 insertions, 4 deletions

diff --git a/include/llvm/CodeGen/MachineBasicBlock.h b/include/llvm/CodeGen/MachineBasicBlock.h
index c917bd8..6cacd69 100644
--- a/include/llvm/CodeGen/MachineBasicBlock.h
+++ b/include/llvm/CodeGen/MachineBasicBlock.h
@@ -572,7 +572,7 @@ private:
   /// getSuccWeight - Return weight of the edge from this block to MBB. This
   /// method should NOT be called directly, but by using getEdgeWeight method
   /// from MachineBranchProbabilityInfo class.
-  uint32_t getSuccWeight(const MachineBasicBlock *succ) const;
+  uint32_t getSuccWeight(const_succ_iterator Succ) const;
 
 
   // Methods used to maintain doubly linked list of blocks...
diff --git a/include/llvm/CodeGen/MachineBranchProbabilityInfo.h b/include/llvm/CodeGen/MachineBranchProbabilityInfo.h
index af4db7d..12189ce 100644
--- a/include/llvm/CodeGen/MachineBranchProbabilityInfo.h
+++ b/include/llvm/CodeGen/MachineBranchProbabilityInfo.h
@@ -16,14 +16,12 @@
 #define LLVM_CODEGEN_MACHINEBRANCHPROBABILITYINFO_H
 
 #include "llvm/Pass.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/Support/BranchProbability.h"
 #include <climits>
 
 namespace llvm {
 
-class raw_ostream;
-class MachineBasicBlock;
-
 class MachineBranchProbabilityInfo : public ImmutablePass {
   virtual void anchor();
 
@@ -52,6 +50,11 @@ public:
   uint32_t getEdgeWeight(const MachineBasicBlock *Src,
                          const MachineBasicBlock *Dst) const;
 
+  // Same thing, but using a const_succ_iterator from Src. This is faster when
+  // the iterator is already available.
+  uint32_t getEdgeWeight(const MachineBasicBlock *Src,
+                         MachineBasicBlock::const_succ_iterator Dst) const;
+
   // Get sum of the block successors' weights, potentially scaling them to fit
   // within 32-bits. If scaling is required, sets Scale based on the necessary
   // adjustment. Any edge weights used with the sum should be divided by Scale.