[llvm] [Coverage] Improve performance of propagating Counter of Expansions (PR #122589)

[NAKAMURA Takumi via llvm-commits]

https://github.com/chapuni created https://github.com/llvm/llvm-project/pull/122589

This improves `clang++ -dump-coverage-mapping` (939,498 lines) for `RISCVInstructionSelector.cpp` from 30m to 1m18s and also improves `llvm-cov` for `check-llvm` from 33m to 24s.

The current implementation behaved O(N^2) order with hundreds thousands of Expansions.

This assumes:
  - Records are partitioned by FileID.
    - ExpandedFileID doesn't point FileID==0, since it is the root.
  - The Count in Expansion is propagated from 1st Record in ExpandedFileID.

Therefore another fact below can be assumed.
  - Propagation chain consists of Expansions at each 1st Record.

This scans the Record at most a few times. O(N) is expected.

Fixes #113173

>From e5fbf846113591b0abff151d91b01a5dd40abef1 Mon Sep 17 00:00:00 2001
From: NAKAMURA Takumi <geek4civic at gmail.com>
Date: Sat, 11 Jan 2025 18:42:10 +0900
Subject: [PATCH] [Coverage] Improve performance of propagating Counter of
 Expansions

This improves `-dump-coverage-mapping` (939,498 lines) for
`RISCVInstructionSelector.cpp` from 30m to 1m18s and also improves
`llvm-cov` for `check-llvm` from 33m to 24s.

The current implementation behaved O(N^2) order with hundreds
thousands of Expansions.

This assumes:
  - Records are partitioned by FileID.
    - ExpandedFileID doesn't point FileID==0, since it is the root.
  - The Count in Expansion is propagated from 1st Record in
    ExpandedFileID.

Therefore another fact below can be assumed.
  - Propagation chain consists of Expansions at each 1st Record.

This scans the Record at most a few times. O(N) is expected.

Fixes #113173
---
 .../Coverage/CoverageMappingReader.cpp        | 100 +++++++++++++++---
 1 file changed, 84 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp b/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
index cdf4412c6477a1..15c7004c45e7d6 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
@@ -440,26 +440,94 @@ Error RawCoverageMappingReader::read() {
   // Set the counters for the expansion regions.
   // i.e. Counter of expansion region = counter of the first region
   // from the expanded file.
-  // Perform multiple passes to correctly propagate the counters through
-  // all the nested expansion regions.
-  SmallVector<CounterMappingRegion *, 8> FileIDExpansionRegionMapping;
-  FileIDExpansionRegionMapping.resize(VirtualFileMapping.size(), nullptr);
-  for (unsigned Pass = 1, S = VirtualFileMapping.size(); Pass < S; ++Pass) {
-    for (auto &R : MappingRegions) {
-      if (R.Kind != CounterMappingRegion::ExpansionRegion)
-        continue;
-      assert(!FileIDExpansionRegionMapping[R.ExpandedFileID]);
-      FileIDExpansionRegionMapping[R.ExpandedFileID] = &R;
+  // This assumes:
+  //   - Records are partitioned by FileID.
+  //     - ExpandedFileID doesn't point FileID==0, since it is the root.
+  //   - The Count in Expansion is propagated from 1st Record in
+  //     ExpandedFileID.
+  // Therefore another fact below can be assumed.
+  //   - Propagation chain consists of Expansions at each 1st Record.
+
+  /// Node per File.
+  struct FileNode {
+    /// 1st Record in the File.
+    CounterMappingRegion *FirstR = nullptr;
+    /// The Expansion Record out of this File.
+    CounterMappingRegion *ExpanderR = nullptr;
+    /// Count hasn't been propagated yet (for Expansion)
+    bool ExpansionPending = false;
+  };
+  SmallVector<FileNode> FileIDExpansionRegionMapping(VirtualFileMapping.size());
+  assert(VirtualFileMapping.size() == FileIDExpansionRegionMapping.size());
+
+  // Construct the tree.
+  unsigned PrevFileID = 0;
+  for (auto &R : MappingRegions) {
+    if (R.Kind == CounterMappingRegion::ExpansionRegion) {
+      // The File that contains Expansion may be a node.
+      assert(R.ExpandedFileID != 0 &&
+             "ExpandedFileID shouldn't point the root");
+      assert(!FileIDExpansionRegionMapping[R.ExpandedFileID].ExpanderR);
+      FileIDExpansionRegionMapping[R.ExpandedFileID].ExpanderR = &R;
+
+      // It will be the node if (isExpansion and ExpansionPending).
+      FileIDExpansionRegionMapping[R.ExpandedFileID].ExpansionPending = true;
     }
-    for (auto &R : MappingRegions) {
-      if (FileIDExpansionRegionMapping[R.FileID]) {
-        FileIDExpansionRegionMapping[R.FileID]->Count = R.Count;
-        FileIDExpansionRegionMapping[R.FileID] = nullptr;
-      }
+
+    // FileID==0 is not handled here but don't care since it's the root.
+    if (PrevFileID != R.FileID) {
+      // Record 1st Record for each File.
+      assert(!FileIDExpansionRegionMapping[R.FileID].FirstR);
+      FileIDExpansionRegionMapping[R.FileID].FirstR = &R;
+      PrevFileID = R.FileID;
     }
   }
 
-  return Error::success();
+  // Make sure the root [0] doesn't point others.
+  assert((FileIDExpansionRegionMapping.empty() ||
+          (!FileIDExpansionRegionMapping[0].ExpanderR &&
+           !FileIDExpansionRegionMapping[0].FirstR)) &&
+         "The root shouldn't point other nodes");
+
+  auto Propagate = [&](FileNode &X) {
+    // Skip already processed node.
+    if (!X.ExpanderR)
+      return false;
+    // Skip Pending Expansion node.
+    // Process non-Expansion Record and non-Pending Expansion.
+    if (X.ExpansionPending &&
+        X.FirstR->Kind == CounterMappingRegion::ExpansionRegion)
+      return false;
+
+    // Propagate.
+    X.ExpanderR->Count = X.FirstR->Count;
+    // Mark the destination ready.
+    FileIDExpansionRegionMapping[X.ExpanderR->FileID].ExpansionPending = false;
+    // Mark this processed.
+    X.ExpanderR = nullptr;
+    return true;
+  };
+
+  // This won't iterate in most cases but iterates finitely for
+  // unusual cases.
+  for (unsigned I = 0, E = FileIDExpansionRegionMapping.size(); I != E; ++I) {
+    // Propagate Descending.  All Files will be processed if each
+    // ExpandedFileID is greater than FileID.
+    for (auto &X : reverse(FileIDExpansionRegionMapping))
+      Propagate(X);
+
+    // Check whether all Files are processed (or propagate Ascending,
+    // not possible in the usual assumptions).
+    bool Changed = false;
+    for (auto &X : FileIDExpansionRegionMapping)
+      Changed = (Propagate(X) || Changed);
+
+    if (!Changed)
+      return Error::success();
+  }
+
+  return make_error<CoverageMapError>(coveragemap_error::malformed,
+                                      "unexpected Expansions");
 }
 
 Expected<bool> RawCoverageMappingDummyChecker::isDummy() {