From e2a44f049b3db137c2d3b5cca338be359bf37167 Mon Sep 17 00:00:00 2001
From: Mike Hommey <mh@glandium.org>
Date: Wed, 8 Jan 2025 13:12:03 +0900
Subject: [PATCH] [Coverage] Speed up function record iteration

When iterating over function records, filtered by file name, currently,
the iteration goes over all the function records, repeatedly for each
source file, essentially giving quadratic behavior.

413647d730972eac9675f695c2ea63fb393a5531 sped up some cases by keeping
track of the indices of the function records corresponding to each
file name. This change expands the use of that map to
FunctionRecordIterator.

On a test case with Firefox's libxul.so and a 2.5MB profile, this brings
down the runtime of `llvm-cov export $lib --instr-profile $prof -t lcov`
from 12 minutes with 90% spent in skipOtherFiles to 19 seconds with no
samples in skipOtherFiles at all under a sampling profiler (with a
sampling interval of 1ms).

Fixes #62079
---
 .../ProfileData/Coverage/CoverageMapping.h    | 46 ++++++++++++++++---
 .../ProfileData/Coverage/CoverageMapping.cpp  |  2 +-
 2 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h b/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h
index 81307d7b025d9..5a20a9ef63287 100644
--- a/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h
+++ b/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h
@@ -748,10 +748,15 @@ struct FunctionRecord {
 };
 
 /// Iterator over Functions, optionally filtered to a single file.
+/// When filtering to a single file, the iterator requires a list of potential
+/// indices where to find the desired records to avoid quadratic behavior when
+/// repeatedly iterating over functions from different files.
 class FunctionRecordIterator
     : public iterator_facade_base<FunctionRecordIterator,
                                   std::forward_iterator_tag, FunctionRecord> {
   ArrayRef<FunctionRecord> Records;
+  ArrayRef<unsigned> RecordIndices;
+  ArrayRef<unsigned>::iterator CurrentIndex;
   ArrayRef<FunctionRecord>::iterator Current;
   StringRef Filename;
 
@@ -760,8 +765,17 @@ class FunctionRecordIterator
 
 public:
   FunctionRecordIterator(ArrayRef<FunctionRecord> Records_,
-                         StringRef Filename = "")
-      : Records(Records_), Current(Records.begin()), Filename(Filename) {
+                         StringRef Filename = "",
+                         ArrayRef<unsigned> RecordIndices_ = {})
+      : Records(Records_), RecordIndices(RecordIndices_),
+        CurrentIndex(RecordIndices.begin()),
+        // If `RecordIndices` is provided, we can skip directly to the first
+        // index it provides.
+        Current(CurrentIndex == RecordIndices.end() ? Records.begin()
+                                                    : &Records[*CurrentIndex]),
+        Filename(Filename) {
+    assert(Filename.empty() == RecordIndices_.empty() &&
+           "If `Filename` is specified, `RecordIndices` must also be provided");
     skipOtherFiles();
   }
 
@@ -774,11 +788,29 @@ class FunctionRecordIterator
   const FunctionRecord &operator*() const { return *Current; }
 
   FunctionRecordIterator &operator++() {
-    assert(Current != Records.end() && "incremented past end");
-    ++Current;
+    advanceOne();
     skipOtherFiles();
     return *this;
   }
+
+private:
+  void advanceOne() {
+    if (RecordIndices.empty()) {
+      // Iteration over all entries, advance in the list of records.
+      assert(Current != Records.end() && "incremented past end");
+      ++Current;
+    } else {
+      // Iterator over entries filtered by file name. Advance in the list of
+      // indices, and adjust the cursor in the list of records accordingly.
+      assert(CurrentIndex != RecordIndices.end() && "incremented past end");
+      ++CurrentIndex;
+      if (CurrentIndex == RecordIndices.end()) {
+        Current = Records.end();
+      } else {
+        Current = &Records[*CurrentIndex];
+      }
+    }
+  }
 };
 
 /// Coverage information for a macro expansion or #included file.
@@ -1037,8 +1069,10 @@ class CoverageMapping {
   /// Gets all of the functions in a particular file.
   iterator_range<FunctionRecordIterator>
   getCoveredFunctions(StringRef Filename) const {
-    return make_range(FunctionRecordIterator(Functions, Filename),
-                      FunctionRecordIterator());
+    return make_range(
+        FunctionRecordIterator(Functions, Filename,
+                               getImpreciseRecordIndicesForFilename(Filename)),
+        FunctionRecordIterator());
   }
 
   /// Get the list of function instantiation groups in a particular file.
diff --git a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
index 6d6678e9e4afe..c39585681911a 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
@@ -618,7 +618,7 @@ unsigned CounterMappingContext::getMaxCounterID(const Counter &C) const {
 void FunctionRecordIterator::skipOtherFiles() {
   while (Current != Records.end() && !Filename.empty() &&
          Filename != Current->Filenames[0])
-    ++Current;
+    advanceOne();
   if (Current == Records.end())
     *this = FunctionRecordIterator();
 }