From 90b0516053169ee91503622398c98f1c747be793 Mon Sep 17 00:00:00 2001
From: Zalathar <Zalathar@users.noreply.github.com>
Date: Sun, 3 Sep 2023 15:52:49 +1000
Subject: [PATCH] coverage: Simplify grouping of mappings by file

This removes an ad-hoc implementation of `group_by`.
---
 .../src/coverageinfo/mapgen.rs                | 46 +++++++++----------
 compiler/rustc_codegen_llvm/src/lib.rs        |  1 +
 2 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
index a8380c6dcb967..3ec83c4f18a8d 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
@@ -165,32 +165,32 @@ fn encode_mappings_for_function(
 
     let mut virtual_file_mapping = Vec::new();
     let mut mapping_regions = Vec::with_capacity(counter_regions.len());
-    let mut current_file_name = None;
-    let mut current_file_id = 0;
-
-    // Convert the list of (Counter, CodeRegion) pairs to an array of `CounterMappingRegion`, sorted
-    // by filename and position. Capture any new files to compute the `CounterMappingRegion`s
-    // `file_id` (indexing files referenced by the current function), and construct the
-    // function-specific `virtual_file_mapping` from `file_id` to its index in the module's
-    // `filenames` array.
+
+    // Sort the list of (counter, region) mapping pairs by region, so that they
+    // can be grouped by filename. Prepare file IDs for each filename, and
+    // prepare the mapping data so that we can pass it through FFI to LLVM.
     counter_regions.sort_by_key(|(_counter, region)| *region);
-    for (counter, region) in counter_regions {
-        let CodeRegion { file_name, start_line, start_col, end_line, end_col } = *region;
-        let same_file = current_file_name.is_some_and(|p| p == file_name);
-        if !same_file {
-            if current_file_name.is_some() {
-                current_file_id += 1;
-            }
-            current_file_name = Some(file_name);
-            debug!("  file_id: {} = '{:?}'", current_file_id, file_name);
-            let global_file_id = global_file_table.global_file_id_for_file_name(file_name);
-            virtual_file_mapping.push(global_file_id);
-        }
-        {
-            debug!("Adding counter {:?} to map for {:?}", counter, region);
+    for (local_file_id, counter_regions_for_file) in
+        (0u32..).zip(counter_regions.group_by(|(_, a), (_, b)| a.file_name == b.file_name))
+    {
+        // Look up (or allocate) the global file ID for this filename.
+        let file_name = counter_regions_for_file[0].1.file_name;
+        let global_file_id = global_file_table.global_file_id_for_file_name(file_name);
+
+        // Associate that global file ID with a local file ID for this function.
+        assert_eq!(local_file_id as usize, virtual_file_mapping.len());
+        virtual_file_mapping.push(global_file_id);
+        debug!("  file_id: local {local_file_id} => global {global_file_id} = '{file_name:?}'");
+
+        // For each counter/region pair in this function+file, convert it to a
+        // form suitable for FFI.
+        for &(counter, region) in counter_regions_for_file {
+            let CodeRegion { file_name: _, start_line, start_col, end_line, end_col } = *region;
+
+            debug!("Adding counter {counter:?} to map for {region:?}");
             mapping_regions.push(CounterMappingRegion::code_region(
                 counter,
-                current_file_id,
+                local_file_id,
                 start_line,
                 start_col,
                 end_line,
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
index d283299ac46b4..ac199624e347b 100644
--- a/compiler/rustc_codegen_llvm/src/lib.rs
+++ b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -10,6 +10,7 @@
 #![feature(iter_intersperse)]
 #![feature(let_chains)]
 #![feature(never_type)]
+#![feature(slice_group_by)]
 #![feature(impl_trait_in_assoc_type)]
 #![recursion_limit = "256"]
 #![allow(rustc::potential_query_instability)]