2 files changed, 73 insertions, 5 deletions

diff --git a/llvm/lib/Analysis/ProfileInfo.cpp b/llvm/lib/Analysis/ProfileInfo.cpp
index a96dc4f2fc6..8cb1e001764 100644
--- a/llvm/lib/Analysis/ProfileInfo.cpp
+++ b/llvm/lib/Analysis/ProfileInfo.cpp
@@ -14,6 +14,8 @@
 
 #include "llvm/Analysis/ProfileInfo.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/CFG.h"
+#include <set>
 using namespace llvm;
 
 // Register the ProfileInfo interface, providing a nice name to refer to.
@@ -23,6 +25,56 @@ namespace {
 
 ProfileInfo::~ProfileInfo() {}
 
+unsigned ProfileInfo::getExecutionCount(BasicBlock *BB) const {
+  pred_iterator PI = pred_begin(BB), PE = pred_end(BB);
+
+  // Are there zero predecessors of this block?
+  if (PI == PE) {
+    // If this is the entry block, look for the Null -> Entry edge.
+    if (BB == &BB->getParent()->front())
+      return getEdgeWeight(0, BB);
+    else
+      return 0;   // Otherwise, this is a dead block.
+  }
+
+  // Otherwise, if there are predecessors, the execution count of this block is
+  // the sum of the edge frequencies from the incoming edges.  Note that if
+  // there are multiple edges from a predecessor to this block that we don't
+  // want to count its weight multiple times.  For this reason, we keep track of
+  // the predecessors we've seen and only count them if we haven't run into them
+  // yet.
+  //
+  // We don't want to create an std::set unless we are dealing with a block that
+  // has a LARGE number of in-edges.  Handle the common case of having only a
+  // few in-edges with special code.
+  //
+  BasicBlock *FirstPred = *PI;
+  unsigned Count = getEdgeWeight(FirstPred, BB);
+  ++PI;
+  if (PI == PE) return Count;   // Quick exit for single predecessor blocks
+
+  BasicBlock *SecondPred = *PI;
+  if (SecondPred != FirstPred) Count += getEdgeWeight(SecondPred, BB);
+  ++PI;
+  if (PI == PE) return Count;   // Quick exit for two predecessor blocks
+
+  BasicBlock *ThirdPred = *PI;
+  if (ThirdPred != FirstPred && ThirdPred != SecondPred)
+    Count += getEdgeWeight(ThirdPred, BB);
+  ++PI;
+  if (PI == PE) return Count;   // Quick exit for three predecessor blocks
+
+  std::set<BasicBlock*> ProcessedPreds;
+  ProcessedPreds.insert(FirstPred);
+  ProcessedPreds.insert(SecondPred);
+  ProcessedPreds.insert(ThirdPred);
+  for (; PI != PE; ++PI)
+    if (ProcessedPreds.insert(*PI).second)
+      Count += getEdgeWeight(*PI, BB);
+  return Count;
+}
+
+
 
 //===----------------------------------------------------------------------===//
 //  NoProfile ProfileInfo implementation
diff --git a/llvm/lib/Analysis/ProfileInfoLoaderPass.cpp b/llvm/lib/Analysis/ProfileInfoLoaderPass.cpp
index c1ff8cf5631..b32dcc406ad 100644
--- a/llvm/lib/Analysis/ProfileInfoLoaderPass.cpp
+++ b/llvm/lib/Analysis/ProfileInfoLoaderPass.cpp
@@ -12,6 +12,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/BasicBlock.h"
+#include "llvm/InstrTypes.h"
 #include "llvm/Pass.h"
 #include "llvm/Analysis/ProfileInfo.h"
 #include "llvm/Analysis/ProfileInfoLoader.h"
@@ -60,11 +62,25 @@ Pass *llvm::createProfileLoaderPass(const std::string &Filename) {
 
 bool LoaderPass::run(Module &M) {
   ProfileInfoLoader PIL("profile-loader", Filename, M);
-  ExecutionCounts.clear();
-  if (PIL.hasAccurateBlockCounts()) {
-    std::vector<std::pair<BasicBlock*, unsigned> > Counts;
-    PIL.getBlockCounts(Counts);
-    ExecutionCounts.insert(Counts.begin(), Counts.end());
+  EdgeCounts.clear();
+  bool PrintedWarning = false;
+  
+  std::vector<std::pair<ProfileInfoLoader::Edge, unsigned> > ECs;
+  PIL.getEdgeCounts(ECs);
+  for (unsigned i = 0, e = ECs.size(); i != e; ++i) {
+    BasicBlock *BB = ECs[i].first.first;
+    unsigned SuccNum = ECs[i].first.second;
+    TerminatorInst *TI = BB->getTerminator();
+    if (SuccNum >= TI->getNumSuccessors()) {
+      if (!PrintedWarning) {
+        std::cerr << "WARNING: profile information is inconsistent with "
+                  << "the current program!\n";
+        PrintedWarning = true;
+      }
+    } else {
+      EdgeCounts[std::make_pair(BB, TI->getSuccessor(SuccNum))]+= ECs[i].second;
+    }
   }
+
   return false;
 }