coverage: Collect a function's coverage mappings into a single list

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -1,8 +1,9 @@
 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::mir::coverage::{CodeRegion, CounterId, CovTerm, ExpressionId, Mapping, Op};
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::TyCtxt;
 
@@ -11,28 +12,21 @@ 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
-/// expressions (additions or subtraction), and (c) unreachable regions (always counted as zero),
-/// for a given Function. This struct also stores the `function_source_hash`,
-/// computed during instrumentation, and forwarded with counters.
-///
-/// Note, it may be important to understand LLVM's definitions of `unreachable` regions versus "gap
-/// regions" (or "gap areas"). A gap region is a code region within a counted region (either counter
-/// or expression), but the line or lines in the gap region are not executable (such as lines with
-/// only whitespace or comments). According to LLVM Code Coverage Mapping documentation, "A count
-/// for a gap area is only used as the line execution count if there are no other regions on a
-/// line."
+/// Holds all of the coverage mapping data associated with a function instance,
+/// collected during traversal of `Coverage` statements in the function's MIR.
 #[derive(Debug)]
 pub struct FunctionCoverage<'tcx> {
     instance: Instance<'tcx>,
     source_hash: u64,
     is_used: bool,
-    counters: IndexVec<CounterId, Option<Vec<CodeRegion>>>,
+
+    /// Tracks which counters have been seen, to avoid duplicate mappings
+    /// that might be introduced by MIR inlining.
+    counters_seen: BitSet<CounterId>,
     expressions: IndexVec<ExpressionId, Option<Expression>>,
-    unreachable_regions: Vec<CodeRegion>,
+    mappings: Vec<Mapping>,
 }
 
 impl<'tcx> FunctionCoverage<'tcx> {
@@ -56,9 +50,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,19 +74,8 @@ 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() {
-            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"
-            ),
+        if self.counters_seen.insert(id) {
+            self.add_mappings(CovTerm::Counter(id), code_regions);
         }
     }
 
@@ -120,10 +103,13 @@ 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),
+            None => {
+                *slot = Some(expression);
+                self.add_mappings(CovTerm::Expression(expression_id), code_regions);
+            }
             // If this expression ID slot has already been filled, it should
             // contain identical information.
             Some(ref previous_expression) => assert_eq!(
@@ -137,7 +123,22 @@ impl<'tcx> FunctionCoverage<'tcx> {
     #[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);
+    }
+
+    #[instrument(level = "debug", skip(self))]
+    fn add_mappings(&mut self, kind: CovTerm, code_regions: &[CodeRegion]) {
+        self.mappings
+            .extend(code_regions.iter().cloned().map(|code_region| Mapping { kind, code_region }));
+    }
+
+    pub(crate) fn finalize(&mut self) {
+        self.assert_source_hash_is_set();
+
+        self.simplify_expressions();
+
+        // Sort all of the collected mappings into a predictable order.
+        self.mappings.sort_unstable();
     }
 
     /// Perform some simplifications to make the final coverage mappings
@@ -146,7 +147,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
@@ -198,49 +199,17 @@ impl<'tcx> FunctionCoverage<'tcx> {
         self.source_hash
     }
 
-    /// Generate an array of CounterExpressions, and an iterator over all `Counter`s and their
-    /// associated `Regions` (from which the LLVM-specific `CoverageMapGenerator` will create
-    /// `CounterMappingRegion`s.
-    pub fn get_expressions_and_counter_regions(
-        &self,
-    ) -> (Vec<CounterExpression>, impl Iterator<Item = (Counter, &CodeRegion)>) {
+    fn assert_source_hash_is_set(&self) {
         assert!(
             self.source_hash != 0 || !self.is_used,
             "No counters provided the source_hash for used function: {:?}",
             self.instance
         );
-
-        let counter_expressions = self.counter_expressions();
-        // Expression IDs are indices into `self.expressions`, and on the LLVM
-        // side they will be treated as indices into `counter_expressions`, so
-        // 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 =
-            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> {
+    pub(crate) fn expressions_for_ffi(&self) -> Vec<CounterExpression> {
         // We know that LLVM will optimize out any unused expressions before
         // producing the final coverage map, so there's no need to do the same
         // thing on the Rust side unless we're confident we can do much better.
@@ -268,24 +237,12 @@ 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))
+    /// Converts this function's coverage mappings into an intermediate form
+    /// that will be used by `mapgen` when preparing for FFI.
+    pub(crate) fn mappings_for_ffi(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
+        self.mappings.iter().map(|&Mapping { kind, ref code_region }| {
+            let counter = Counter::from_term(kind);
+            (counter, code_region)
+        })
     }
 }

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -62,7 +62,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();
@@ -159,20 +160,21 @@ fn encode_mappings_for_function(
     global_file_table: &mut GlobalFileTable,
     function_coverage: &FunctionCoverage<'_>,
 ) -> Vec<u8> {
-    let (expressions, counter_regions) = function_coverage.get_expressions_and_counter_regions();
+    let expressions = function_coverage.expressions_for_ffi();
+    let mut counter_regions = function_coverage.mappings_for_ffi().collect::<Vec<_>>();
 
-    let mut counter_regions = counter_regions.collect::<Vec<_>>();
     if counter_regions.is_empty() {
         return Vec::new();
     }
 
     let mut virtual_file_mapping = IndexVec::<u32, u32>::new();
     let mut mapping_regions = Vec::with_capacity(counter_regions.len());
 
-    // 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);
+    // Sort and group the list of (counter, region) mapping pairs by filename.
+    // (Preserve any further ordering imposed by `FunctionCoverage`.)
+    // 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.file_name);
     for counter_regions_for_file in
         counter_regions.group_by(|(_, a), (_, b)| a.file_name == b.file_name)
     {

compiler/rustc_middle/src/mir/coverage.rs

@@ -35,18 +35,19 @@ impl ExpressionId {
     pub const START: Self = Self::from_u32(0);
 }
 
-/// Enum that can hold a constant zero value, the ID of an physical coverage
-/// counter, or the ID of a coverage-counter expression.
+/// Enum that can hold the ID of a physical coverage counter, the ID of a
+/// coverage-counter expression, or a constant zero value.
 ///
 /// This was originally only used for expression operands (and named `Operand`),
 /// but the zero/counter/expression distinction is also useful for representing
 /// the value of code/gap mappings, and the true/false arms of branch mappings.
-#[derive(Copy, Clone, PartialEq, Eq)]
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub enum CovTerm {
-    Zero,
+    // Coverage codegen relies on this variant order when it sorts a function's mappings.
     Counter(CounterId),
     Expression(ExpressionId),
+    Zero,
 }
 
 impl Debug for CovTerm {
@@ -102,6 +103,8 @@ impl Debug for CoverageKind {
 #[derive(Clone, TyEncodable, TyDecodable, Hash, HashStable, PartialEq, Eq, PartialOrd, Ord)]
 #[derive(TypeFoldable, TypeVisitable)]
 pub struct CodeRegion {
+    // Coverage codegen relies on filenames being sorted first, because it
+    // needs to group mappings by file.
     pub file_name: Symbol,
     pub start_line: u32,
     pub start_col: u32,
@@ -135,3 +138,17 @@ impl Op {
         matches!(self, Self::Subtract)
     }
 }
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Mapping {
+    // Coverage codegen relies on this field order when it sorts a function's mappings.
+    pub code_region: CodeRegion,
+
+    /// Kind of coverage region that this mapping represents.
+    ///
+    /// Currently we only support ordinary `Code` regions, so the "kind" is
+    /// just a term. When we add support for other region kinds allowed by
+    /// LLVM (e.g. branch regions, expansion regions), we'll need to replace
+    /// this with a dedicated mapping-kind enum.
+    pub kind: CovTerm,
+}