1 files changed, 77 insertions, 23 deletions

diff --git a/llvm/tools/llvm-cov/SourceCoverageDataManager.cpp b/llvm/tools/llvm-cov/SourceCoverageDataManager.cpp
index 6f7c8e5d985..a60ea6e2586 100644
--- a/llvm/tools/llvm-cov/SourceCoverageDataManager.cpp
+++ b/llvm/tools/llvm-cov/SourceCoverageDataManager.cpp
@@ -19,32 +19,86 @@ using namespace coverage;
 
 void SourceCoverageDataManager::insert(const CountedRegion &CR) {
   Regions.push_back(CR);
-  Uniqued = false;
+  Segments.clear();
 }
 
-ArrayRef<CountedRegion> SourceCoverageDataManager::getSourceRegions() {
-  if (Uniqued || Regions.size() <= 1)
-    return Regions;
-
-  // Sort the regions given that they're all in the same file at this point.
-  std::sort(Regions.begin(), Regions.end(),
-            [](const CountedRegion &LHS, const CountedRegion &RHS) {
-    return LHS.startLoc() < RHS.startLoc();
-  });
-
-  // Merge duplicate source ranges and sum their execution counts.
-  auto Prev = Regions.begin();
-  for (auto I = Prev + 1, E = Regions.end(); I != E; ++I) {
-    if (I->startLoc() == Prev->startLoc() && I->endLoc() == Prev->endLoc()) {
-      Prev->ExecutionCount += I->ExecutionCount;
-      continue;
+namespace {
+class SegmentBuilder {
+  std::vector<CoverageSegment> Segments;
+  SmallVector<const CountedRegion *, 8> ActiveRegions;
+
+  /// Start a segment with no count specified.
+  void startSegment(unsigned Line, unsigned Col, bool IsRegionEntry) {
+    Segments.emplace_back(Line, Col, IsRegionEntry);
+  }
+
+  /// Start a segment with the given Region's count.
+  void startSegment(unsigned Line, unsigned Col, bool IsRegionEntry,
+                    const CountedRegion &Region) {
+    if (Segments.empty())
+      Segments.emplace_back(Line, Col, IsRegionEntry);
+    CoverageSegment S = Segments.back();
+    // Avoid creating empty regions.
+    if (S.Line != Line || S.Col != Col) {
+      Segments.emplace_back(Line, Col, IsRegionEntry);
+      S = Segments.back();
+    }
+    // Set this region's count.
+    if (Region.Kind != coverage::CounterMappingRegion::SkippedRegion)
+      Segments.back().setCount(Region.ExecutionCount);
+  }
+
+  /// Start a segment for the given region.
+  void startSegment(const CountedRegion &Region) {
+    startSegment(Region.LineStart, Region.ColumnStart, true, Region);
+  }
+
+  /// Pop the top region off of the active stack, starting a new segment with
+  /// the containing Region's count.
+  void popRegion() {
+    const CountedRegion *Active = ActiveRegions.back();
+    unsigned Line = Active->LineEnd, Col = Active->ColumnEnd;
+    ActiveRegions.pop_back();
+    if (ActiveRegions.empty())
+      startSegment(Line, Col, /*IsRegionEntry=*/false);
+    else
+      startSegment(Line, Col, /*IsRegionEntry=*/false, *ActiveRegions.back());
+  }
+
+public:
+  /// Build a list of CoverageSegments from a sorted list of Regions.
+  std::vector<CoverageSegment>
+  buildSegments(ArrayRef<CountedRegion> Regions) {
+    for (const auto &Region : Regions) {
+      // Pop any regions that end before this one starts.
+      while (!ActiveRegions.empty() &&
+             ActiveRegions.back()->endLoc() <= Region.startLoc())
+        popRegion();
+      // Add this region to the stack.
+      ActiveRegions.push_back(&Region);
+      startSegment(Region);
     }
-    ++Prev;
-    *Prev = *I;
+    // Pop any regions that are left in the stack.
+    while (!ActiveRegions.empty())
+      popRegion();
+    return Segments;
+  }
+};
+}
+
+ArrayRef<CoverageSegment> SourceCoverageDataManager::getCoverageSegments() {
+  if (Segments.empty()) {
+    // Sort the regions given that they're all in the same file at this point.
+    std::sort(Regions.begin(), Regions.end(),
+              [](const CountedRegion &LHS, const CountedRegion &RHS) {
+      if (LHS.startLoc() == RHS.startLoc())
+        // When LHS completely contains RHS, we sort LHS first.
+        return RHS.endLoc() < LHS.endLoc();
+      return LHS.startLoc() < RHS.startLoc();
+    });
+
+    Segments = SegmentBuilder().buildSegments(Regions);
   }
-  ++Prev;
-  Regions.erase(Prev, Regions.end());
 
-  Uniqued = true;
-  return Regions;
+  return Segments;
 }