!! (WIP) Store a function's mappings in a single list

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -1,6 +1,7 @@
 use crate::coverageinfo::ffi::{Counter, CounterExpression, ExprKind};
 
 use rustc_data_structures::fx::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;
@@ -11,7 +12,6 @@ pub struct Expression {
     lhs: CovTerm,
     op: Op,
     rhs: CovTerm,
-    code_regions: Vec<CodeRegion>,
 }
 
 /// Collects all of the coverage regions associated with (a) injected counters, (b) counter
@@ -30,9 +30,10 @@ pub struct FunctionCoverage<'tcx> {
     instance: Instance<'tcx>,
     source_hash: u64,
     is_used: bool,
-    counters: IndexVec<CounterId, Option<Vec<CodeRegion>>>,
+
+    counters_seen: BitSet<CounterId>,
     expressions: IndexVec<ExpressionId, Option<Expression>>,
-    unreachable_regions: Vec<CodeRegion>,
+    mappings: Vec<(CovTerm, CodeRegion)>,
 }
 
 impl<'tcx> FunctionCoverage<'tcx> {
@@ -56,9 +57,9 @@ impl<'tcx> FunctionCoverage<'tcx> {
             instance,
             source_hash: 0, // will be set with the first `add_counter()`
             is_used,
-            counters: IndexVec::from_elem_n(None, coverageinfo.num_counters as usize),
+            counters_seen: BitSet::new_empty(coverageinfo.num_counters as usize),
             expressions: IndexVec::from_elem_n(None, coverageinfo.num_expressions as usize),
-            unreachable_regions: Vec::new(),
+            mappings: Vec::new(),
         }
     }
 
@@ -80,20 +81,10 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// Adds code regions to be counted by an injected counter intrinsic.
     #[instrument(level = "debug", skip(self))]
     pub(crate) fn add_counter(&mut self, id: CounterId, code_regions: &[CodeRegion]) {
-        if code_regions.is_empty() {
+        if !self.counters_seen.insert(id) {
             return;
         }
-
-        let slot = &mut self.counters[id];
-        match slot {
-            None => *slot = Some(code_regions.to_owned()),
-            // If this counter ID slot has already been filled, it should
-            // contain identical information.
-            Some(ref previous_regions) => assert_eq!(
-                previous_regions, code_regions,
-                "add_counter: code regions for id changed"
-            ),
-        }
+        self.add_mappings(CovTerm::Counter(id), code_regions);
     }
 
     /// Adds information about a coverage expression, along with zero or more
@@ -120,24 +111,45 @@ impl<'tcx> FunctionCoverage<'tcx> {
             self,
         );
 
-        let expression = Expression { lhs, op, rhs, code_regions: code_regions.to_owned() };
+        let expression = Expression { lhs, op, rhs };
         let slot = &mut self.expressions[expression_id];
         match slot {
             None => *slot = Some(expression),
             // If this expression ID slot has already been filled, it should
             // contain identical information.
-            Some(ref previous_expression) => assert_eq!(
-                previous_expression, &expression,
-                "add_counter_expression: expression for id changed"
-            ),
+            Some(ref previous_expression) => {
+                assert_eq!(
+                    previous_expression, &expression,
+                    "add_counter_expression: expression for id changed"
+                );
+                return;
+            }
         }
+
+        self.add_mappings(CovTerm::Expression(expression_id), code_regions);
     }
 
     /// Adds regions that will be marked as "unreachable", with a constant "zero counter".
     #[instrument(level = "debug", skip(self))]
     pub(crate) fn add_unreachable_regions(&mut self, code_regions: &[CodeRegion]) {
         assert!(!code_regions.is_empty(), "unreachable regions always have code regions");
-        self.unreachable_regions.extend_from_slice(code_regions);
+        self.add_mappings(CovTerm::Zero, code_regions);
+    }
+
+    fn add_mappings(&mut self, kind: CovTerm, code_regions: &[CodeRegion]) {
+        self.mappings
+            .extend(code_regions.iter().cloned().map(move |code_region| (kind, code_region)));
+    }
+
+    pub(crate) fn finalize(&mut self) {
+        self.simplify_expressions();
+
+        let mapping_kind_sort_key = |mapping_kind: &CovTerm| match mapping_kind {
+            CovTerm::Counter(_) => 0,
+            CovTerm::Expression(_) => 1,
+            CovTerm::Zero => u8::MAX,
+        };
+        self.mappings.sort_by_key(|(mapping_kind, _)| mapping_kind_sort_key(mapping_kind));
     }
 
     /// Perform some simplifications to make the final coverage mappings
@@ -146,7 +158,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// This method mainly exists to preserve the simplifications that were
     /// already being performed by the Rust-side expression renumbering, so that
     /// the resulting coverage mappings don't get worse.
-    pub(crate) fn simplify_expressions(&mut self) {
+    fn simplify_expressions(&mut self) {
         // The set of expressions that either were optimized out entirely, or
         // have zero as both of their operands, and will therefore always have
         // a value of zero. Other expressions that refer to these as operands
@@ -216,28 +228,14 @@ impl<'tcx> FunctionCoverage<'tcx> {
         // the two vectors should correspond 1:1.
         assert_eq!(self.expressions.len(), counter_expressions.len());
 
-        let counter_regions = self.counter_regions();
-        let expression_regions = self.expression_regions();
-        let unreachable_regions = self.unreachable_regions();
+        let counter_regions = self.mappings.iter().map(|(mapping_kind, code_region)| {
+            let counter = Counter::from_term(*mapping_kind);
+            (counter, code_region)
+        });
 
-        let counter_regions =
-            counter_regions.chain(expression_regions.into_iter().chain(unreachable_regions));
         (counter_expressions, counter_regions)
     }
 
-    fn counter_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.counters
-            .iter_enumerated()
-            // Filter out counter IDs that we never saw during MIR traversal.
-            // This can happen if a counter was optimized out by MIR transforms
-            // (and replaced with `CoverageKind::Unreachable` instead).
-            .filter_map(|(id, maybe_code_regions)| Some((id, maybe_code_regions.as_ref()?)))
-            .flat_map(|(id, code_regions)| {
-                let counter = Counter::counter_value_reference(id);
-                code_regions.iter().map(move |region| (counter, region))
-            })
-    }
-
     /// Convert this function's coverage expression data into a form that can be
     /// passed through FFI to LLVM.
     fn counter_expressions(&self) -> Vec<CounterExpression> {
@@ -267,25 +265,4 @@ impl<'tcx> FunctionCoverage<'tcx> {
             })
             .collect::<Vec<_>>()
     }
-
-    fn expression_regions(&self) -> Vec<(Counter, &CodeRegion)> {
-        // Find all of the expression IDs that weren't optimized out AND have
-        // one or more attached code regions, and return the corresponding
-        // mappings as counter/region pairs.
-        self.expressions
-            .iter_enumerated()
-            .filter_map(|(id, maybe_expression)| {
-                let code_regions = &maybe_expression.as_ref()?.code_regions;
-                Some((id, code_regions))
-            })
-            .flat_map(|(id, code_regions)| {
-                let counter = Counter::expression(id);
-                code_regions.iter().map(move |code_region| (counter, code_region))
-            })
-            .collect::<Vec<_>>()
-    }
-
-    fn unreachable_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.unreachable_regions.iter().map(|region| (Counter::ZERO, region))
-    }
 }

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -61,7 +61,8 @@ pub fn finalize(cx: &CodegenCx<'_, '_>) {
     let mut function_data = Vec::new();
     for (instance, mut function_coverage) in function_coverage_map {
         debug!("Generate function coverage for {}, {:?}", cx.codegen_unit.name(), instance);
-        function_coverage.simplify_expressions();
+        function_coverage.finalize();
+        let function_coverage = function_coverage;
 
         let mangled_function_name = tcx.symbol_name(instance).name;
         let source_hash = function_coverage.source_hash();