1 files changed, 15 insertions, 21 deletions

diff --git a/llvm/tools/llvm-cov/SourceCoverageView.cpp b/llvm/tools/llvm-cov/SourceCoverageView.cpp
index f944dc6..337beb8 100644
--- a/llvm/tools/llvm-cov/SourceCoverageView.cpp
+++ b/llvm/tools/llvm-cov/SourceCoverageView.cpp
@@ -187,36 +187,29 @@ void SourceCoverageView::print(raw_ostream &OS, bool WholeFile,
   auto EndISV = InstantiationSubViews.end();
 
   // Get the coverage information for the file.
-  auto NextSegment = CoverageInfo.begin();
+  auto StartSegment = CoverageInfo.begin();
   auto EndSegment = CoverageInfo.end();
+  LineCoverageIterator LCI{CoverageInfo, 1};
+  LineCoverageIterator LCIEnd = LCI.getEnd();
 
-  unsigned FirstLine = NextSegment != EndSegment ? NextSegment->Line : 0;
-  const coverage::CoverageSegment *WrappedSegment = nullptr;
-  SmallVector<const coverage::CoverageSegment *, 8> LineSegments;
-  for (line_iterator LI(File, /*SkipBlanks=*/false); !LI.is_at_eof(); ++LI) {
+  unsigned FirstLine = StartSegment != EndSegment ? StartSegment->Line : 0;
+  for (line_iterator LI(File, /*SkipBlanks=*/false); !LI.is_at_eof();
+       ++LI, ++LCI) {
     // If we aren't rendering the whole file, we need to filter out the prologue
     // and epilogue.
     if (!WholeFile) {
-      if (NextSegment == EndSegment)
+      if (LCI == LCIEnd)
         break;
       else if (LI.line_number() < FirstLine)
         continue;
     }
 
-    // Collect the coverage information relevant to this line.
-    if (LineSegments.size())
-      WrappedSegment = LineSegments.back();
-    LineSegments.clear();
-    while (NextSegment != EndSegment && NextSegment->Line == LI.line_number())
-      LineSegments.push_back(&*NextSegment++);
-
     renderLinePrefix(OS, ViewDepth);
     if (getOptions().ShowLineNumbers)
       renderLineNumberColumn(OS, LI.line_number());
 
-    LineCoverageStats LineCount{LineSegments, WrappedSegment};
     if (getOptions().ShowLineStats)
-      renderLineCoverageColumn(OS, LineCount);
+      renderLineCoverageColumn(OS, *LCI);
 
     // If there are expansion subviews, we want to highlight the first one.
     unsigned ExpansionColumn = 0;
@@ -225,12 +218,12 @@ void SourceCoverageView::print(raw_ostream &OS, bool WholeFile,
       ExpansionColumn = NextESV->getStartCol();
 
     // Display the source code for the current line.
-    renderLine(OS, {*LI, LI.line_number()}, WrappedSegment, LineSegments,
-               ExpansionColumn, ViewDepth);
+    renderLine(OS, {*LI, LI.line_number()}, LCI->getWrappedSegment(),
+               LCI->getLineSegments(), ExpansionColumn, ViewDepth);
 
     // Show the region markers.
-    if (shouldRenderRegionMarkers(LineSegments))
-      renderRegionMarkers(OS, LineSegments, ViewDepth);
+    if (shouldRenderRegionMarkers(LCI->getLineSegments()))
+      renderRegionMarkers(OS, LCI->getLineSegments(), ViewDepth);
 
     // Show the expansions and instantiations for this line.
     bool RenderedSubView = false;
@@ -242,8 +235,9 @@ void SourceCoverageView::print(raw_ostream &OS, bool WholeFile,
       // this subview.
       if (RenderedSubView) {
         ExpansionColumn = NextESV->getStartCol();
-        renderExpansionSite(OS, {*LI, LI.line_number()}, WrappedSegment,
-                            LineSegments, ExpansionColumn, ViewDepth);
+        renderExpansionSite(OS, {*LI, LI.line_number()},
+                            LCI->getWrappedSegment(), LCI->getLineSegments(),
+                            ExpansionColumn, ViewDepth);
         renderViewDivider(OS, ViewDepth + 1);
       }