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

[Jessica Paquette via llvm-commits]

================
@@ -440,26 +441,83 @@ 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.
+  //   - 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());
+
+  // Construct the tree.
+  auto PrevFileID = std::numeric_limits<unsigned>::max(); // Invalid value
+  for (auto &R : MappingRegions) {
+    if (R.Kind == CounterMappingRegion::ExpansionRegion) {
+      // The File that contains Expansion may be a node.
+      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;
-      }
+
+    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();
+  auto Propagate = [&](FileNode &X) {
----------------
ornata wrote:

Why a lambda and not a function?

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