2 files changed, 19 insertions, 11 deletions

diff --git a/include/llvm/Support/GCOV.h b/include/llvm/Support/GCOV.h
index ee81229ff8..c709af5d50 100644
--- a/include/llvm/Support/GCOV.h
+++ b/include/llvm/Support/GCOV.h
@@ -250,6 +250,7 @@ public:
   }
   void addLine(uint32_t N) { Lines.push_back(N); }
   void addCount(size_t DstEdgeNo, uint64_t N);
+  uint64_t getCount() const { return Counter; }
   size_t getNumSrcEdges() const { return SrcEdges.size(); }
   size_t getNumDstEdges() const { return DstEdges.size(); }
 
@@ -269,17 +270,18 @@ private:
   SmallVector<uint32_t, 16> Lines;
 };
 
-typedef DenseMap<uint32_t, uint64_t> LineCounts;
+typedef SmallVector<const GCOVBlock *, 4> BlockVector;
+typedef DenseMap<uint32_t, BlockVector> LineData;
 class FileInfo {
 public:
-  void addLineCount(StringRef Filename, uint32_t Line, uint64_t Count) {
-    LineInfo[Filename][Line-1] += Count;
+  void addBlockLine(StringRef Filename, uint32_t Line, const GCOVBlock *Block) {
+    LineInfo[Filename][Line-1].push_back(Block);
   }
   void setRunCount(uint32_t Runs) { RunCount = Runs; }
   void setProgramCount(uint32_t Programs) { ProgramCount = Programs; }
   void print(raw_fd_ostream &OS, StringRef gcnoFile, StringRef gcdaFile) const;
 private:
-  StringMap<LineCounts> LineInfo;
+  StringMap<LineData> LineInfo;
   uint32_t RunCount;
   uint32_t ProgramCount;
 };
diff --git a/lib/IR/GCOV.cpp b/lib/IR/GCOV.cpp
index ebf10927fc..bcc62b11d5 100644
--- a/lib/IR/GCOV.cpp
+++ b/lib/IR/GCOV.cpp
@@ -282,7 +282,7 @@ void GCOVBlock::addCount(size_t DstEdgeNo, uint64_t N) {
 void GCOVBlock::collectLineCounts(FileInfo &FI) {
   for (SmallVectorImpl<uint32_t>::iterator I = Lines.begin(),
          E = Lines.end(); I != E; ++I)
-    FI.addLineCount(Parent.getFilename(), *I, Counter);
+    FI.addBlockLine(Parent.getFilename(), *I, this);
 }
 
 /// dump - Dump GCOVBlock content to dbgs() for debugging purposes.
@@ -319,7 +319,7 @@ void GCOVBlock::dump() const {
 /// print -  Print source files with collected line count information.
 void FileInfo::print(raw_fd_ostream &OS, StringRef gcnoFile,
                      StringRef gcdaFile) const {
-  for (StringMap<LineCounts>::const_iterator I = LineInfo.begin(),
+  for (StringMap<LineData>::const_iterator I = LineInfo.begin(),
          E = LineInfo.end(); I != E; ++I) {
     StringRef Filename = I->first();
     OwningPtr<MemoryBuffer> Buff;
@@ -335,15 +335,21 @@ void FileInfo::print(raw_fd_ostream &OS, StringRef gcnoFile,
     OS << "        -:    0:Runs:" << RunCount << "\n";
     OS << "        -:    0:Programs:" << ProgramCount << "\n";
 
-    const LineCounts &L = I->second;
+    const LineData &L = I->second;
     uint32_t i = 0;
     while (!AllLines.empty()) {
-      LineCounts::const_iterator CountIt = L.find(i);
-      if (CountIt != L.end()) {
-        if (CountIt->second == 0)
+      LineData::const_iterator BlocksIt = L.find(i);
+      if (BlocksIt != L.end()) {
+        const BlockVector &Blocks = BlocksIt->second;
+        uint64_t LineCount = 0;
+        for (BlockVector::const_iterator I = Blocks.begin(), E = Blocks.end();
+               I != E; ++I) {
+          LineCount += (*I)->getCount();
+        }
+        if (LineCount == 0)
           OS << "    #####:";
         else
-          OS << format("%9" PRIu64 ":", CountIt->second);
+          OS << format("%9" PRIu64 ":", LineCount);
       } else {
         OS << "        -:";
       }