2 files changed, 41 insertions, 22 deletions

diff --git a/llvm/tools/llvm-cov/SourceCoverageView.cpp b/llvm/tools/llvm-cov/SourceCoverageView.cpp
index 52b8ff1747f..6630f3333bd 100644
--- a/llvm/tools/llvm-cov/SourceCoverageView.cpp
+++ b/llvm/tools/llvm-cov/SourceCoverageView.cpp
@@ -83,6 +83,41 @@ CoveragePrinter::create(const CoverageViewOptions &Opts) {
   llvm_unreachable("Unknown coverage output format!");
 }
 
+LineCoverageStats::LineCoverageStats(
+    ArrayRef<const coverage::CoverageSegment *> LineSegments,
+    const coverage::CoverageSegment *WrappedSegment) {
+  // Find the minimum number of regions which start in this line.
+  unsigned MinRegionCount = 0;
+  auto isStartOfRegion = [](const coverage::CoverageSegment *S) {
+    return S->HasCount && S->IsRegionEntry;
+  };
+  for (unsigned I = 0; I < LineSegments.size() && MinRegionCount < 2; ++I)
+    if (isStartOfRegion(LineSegments[I]))
+      ++MinRegionCount;
+
+  ExecutionCount = 0;
+  HasMultipleRegions = MinRegionCount > 1;
+  Mapped = (WrappedSegment && WrappedSegment->HasCount) || (MinRegionCount > 0);
+
+  if (!Mapped)
+    return;
+
+  // Pick the max count among regions which start and end on this line, to
+  // avoid erroneously using the wrapped count, and to avoid picking region
+  // counts which come from deferred regions.
+  if (LineSegments.size() > 1) {
+    for (unsigned I = 0; I < LineSegments.size() - 1; ++I)
+      ExecutionCount = std::max(ExecutionCount, LineSegments[I]->Count);
+    return;
+  }
+
+  // Just pick the maximum count.
+  if (WrappedSegment && WrappedSegment->HasCount)
+    ExecutionCount = WrappedSegment->Count;
+  if (!LineSegments.empty())
+    ExecutionCount = std::max(ExecutionCount, LineSegments[0]->Count);
+}
+
 unsigned SourceCoverageView::getFirstUncoveredLineNo() {
   auto CheckIfUncovered = [](const coverage::CoverageSegment &S) {
     return S.HasCount && S.Count == 0;
@@ -207,17 +242,11 @@ void SourceCoverageView::print(raw_ostream &OS, bool WholeFile,
     while (NextSegment != EndSegment && NextSegment->Line == LI.line_number())
       LineSegments.push_back(&*NextSegment++);
 
-    // Calculate a count to be for the line as a whole.
-    LineCoverageStats LineCount;
-    if (WrappedSegment && WrappedSegment->HasCount)
-      LineCount.addRegionCount(WrappedSegment->Count);
-    for (const auto *S : LineSegments)
-      if (S->HasCount && S->IsRegionEntry)
-        LineCount.addRegionStartCount(S->Count);
-
     renderLinePrefix(OS, ViewDepth);
     if (getOptions().ShowLineNumbers)
       renderLineNumberColumn(OS, LI.line_number());
+
+    LineCoverageStats LineCount{LineSegments, WrappedSegment};
     if (getOptions().ShowLineStats)
       renderLineCoverageColumn(OS, LineCount);
 
diff --git a/llvm/tools/llvm-cov/SourceCoverageView.h b/llvm/tools/llvm-cov/SourceCoverageView.h
index 9cb608fed60..c9f0c57b5cb 100644
--- a/llvm/tools/llvm-cov/SourceCoverageView.h
+++ b/llvm/tools/llvm-cov/SourceCoverageView.h
@@ -67,25 +67,15 @@ struct InstantiationView {
 /// \brief Coverage statistics for a single line.
 struct LineCoverageStats {
   uint64_t ExecutionCount;
-  unsigned RegionCount;
+  bool HasMultipleRegions;
   bool Mapped;
 
-  LineCoverageStats() : ExecutionCount(0), RegionCount(0), Mapped(false) {}
+  LineCoverageStats(ArrayRef<const coverage::CoverageSegment *> LineSegments,
+                    const coverage::CoverageSegment *WrappedSegment);
 
   bool isMapped() const { return Mapped; }
 
-  bool hasMultipleRegions() const { return RegionCount > 1; }
-
-  void addRegionStartCount(uint64_t Count) {
-    // The max of all region starts is the most interesting value.
-    addRegionCount(RegionCount ? std::max(ExecutionCount, Count) : Count);
-    ++RegionCount;
-  }
-
-  void addRegionCount(uint64_t Count) {
-    Mapped = true;
-    ExecutionCount = Count;
-  }
+  bool hasMultipleRegions() const { return HasMultipleRegions; }
 };
 
 /// \brief A file manager that handles format-aware file creation.