[Coverage] Speed up function record iteration

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.

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();
 }