!! (WIP) Store a function's mappings separated by filename

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -1,11 +1,12 @@
 use crate::coverageinfo::ffi::{Counter, CounterExpression, ExprKind};
 
-use rustc_data_structures::fx::FxIndexSet;
+use rustc_data_structures::fx::{FxHashMap, FxIndexSet};
 use rustc_index::bit_set::BitSet;
 use rustc_index::IndexVec;
 use rustc_middle::mir::coverage::{CodeRegion, CounterId, CovTerm, ExpressionId, Op};
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::TyCtxt;
+use rustc_span::Symbol;
 
 #[derive(Clone, Debug, PartialEq)]
 pub struct Expression {
@@ -33,7 +34,10 @@ pub struct FunctionCoverage<'tcx> {
 
     counters_seen: BitSet<CounterId>,
     expressions: IndexVec<ExpressionId, Option<Expression>>,
-    mappings: Vec<(CovTerm, CodeRegion)>,
+
+    /// Invariant: Each code region's filename must match the filename key that
+    /// it is stored under.
+    mappings_per_file: FxHashMap<Symbol, Vec<(CovTerm, CodeRegion)>>,
 }
 
 impl<'tcx> FunctionCoverage<'tcx> {
@@ -59,7 +63,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
             is_used,
             counters_seen: BitSet::new_empty(coverageinfo.num_counters as usize),
             expressions: IndexVec::from_elem_n(None, coverageinfo.num_expressions as usize),
-            mappings: Vec::new(),
+            mappings_per_file: FxHashMap::default(),
         }
     }
 
@@ -137,8 +141,14 @@ impl<'tcx> FunctionCoverage<'tcx> {
     }
 
     fn add_mappings(&mut self, kind: CovTerm, code_regions: &[CodeRegion]) {
-        self.mappings
-            .extend(code_regions.iter().cloned().map(move |code_region| (kind, code_region)));
+        // There's no need to sort the regions by filename before grouping here,
+        // because the result is the same either way.
+        for regions_for_file in code_regions.group_by(|a, b| a.file_name == b.file_name) {
+            let file_name = regions_for_file[0].file_name;
+            self.mappings_per_file.entry(file_name).or_default().extend(
+                regions_for_file.iter().cloned().map(move |code_region| (kind, code_region)),
+            );
+        }
     }
 
     pub(crate) fn finalize(&mut self) {
@@ -149,7 +159,9 @@ impl<'tcx> FunctionCoverage<'tcx> {
             CovTerm::Expression(_) => 1,
             CovTerm::Zero => u8::MAX,
         };
-        self.mappings.sort_by_key(|(mapping_kind, _)| mapping_kind_sort_key(mapping_kind));
+        for mappings in self.mappings_per_file.values_mut() {
+            mappings.sort_by_key(|(mapping_kind, _)| mapping_kind_sort_key(mapping_kind));
+        }
     }
 
     /// Perform some simplifications to make the final coverage mappings
@@ -228,7 +240,8 @@ impl<'tcx> FunctionCoverage<'tcx> {
         // the two vectors should correspond 1:1.
         assert_eq!(self.expressions.len(), counter_expressions.len());
 
-        let counter_regions = self.mappings.iter().map(|(mapping_kind, code_region)| {
+        let all_mappings = self.mappings_per_file.values().flatten();
+        let counter_regions = all_mappings.map(|(mapping_kind, code_region)| {
             let counter = Counter::from_term(*mapping_kind);
             (counter, code_region)
         });