Rollup merge of #119842 - Zalathar:kind, r=oli-obk

coverage: Add enums to accommodate other kinds of coverage mappings

Extracted from  #118305.

LLVM supports several different kinds of coverage mapping regions, but currently we only ever emit ordinary “code” regions.  This PR performs the plumbing required to add other kinds of regions as enum variants, but does not add any specific variants other than `Code`.

The main motivation for this change is branch coverage, but it will also allow separate experimentation with gap regions and skipped regions, which might help in producing more accurate and useful coverage reports.

---

``@rustbot`` label +A-code-coverage

compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs

@@ -1,4 +1,4 @@
-use rustc_middle::mir::coverage::{CounterId, CovTerm, ExpressionId};
+use rustc_middle::mir::coverage::{CodeRegion, CounterId, CovTerm, ExpressionId, MappingKind};
 
 /// Must match the layout of `LLVMRustCounterKind`.
 #[derive(Copy, Clone, Debug)]
@@ -149,6 +149,24 @@ pub struct CounterMappingRegion {
 }
 
 impl CounterMappingRegion {
+    pub(crate) fn from_mapping(
+        mapping_kind: &MappingKind,
+        local_file_id: u32,
+        code_region: &CodeRegion,
+    ) -> Self {
+        let &CodeRegion { file_name: _, start_line, start_col, end_line, end_col } = code_region;
+        match *mapping_kind {
+            MappingKind::Code(term) => Self::code_region(
+                Counter::from_term(term),
+                local_file_id,
+                start_line,
+                start_col,
+                end_line,
+                end_col,
+            ),
+        }
+    }
+
     pub(crate) fn code_region(
         counter: Counter,
         file_id: u32,

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -4,7 +4,8 @@ use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_index::bit_set::BitSet;
 use rustc_middle::mir::coverage::{
-    CodeRegion, CounterId, CovTerm, Expression, ExpressionId, FunctionCoverageInfo, Mapping, Op,
+    CodeRegion, CounterId, CovTerm, Expression, ExpressionId, FunctionCoverageInfo, Mapping,
+    MappingKind, Op,
 };
 use rustc_middle::ty::Instance;
 use rustc_span::Symbol;
@@ -64,8 +65,8 @@ impl<'tcx> FunctionCoverageCollector<'tcx> {
         // For each expression ID that is directly used by one or more mappings,
         // mark it as not-yet-seen. This indicates that we expect to see a
         // corresponding `ExpressionUsed` statement during MIR traversal.
-        for Mapping { term, .. } in &function_coverage_info.mappings {
-            if let &CovTerm::Expression(id) = term {
+        for term in function_coverage_info.mappings.iter().flat_map(|m| m.kind.terms()) {
+            if let CovTerm::Expression(id) = term {
                 expressions_seen.remove(id);
             }
         }
@@ -221,20 +222,21 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// that will be used by `mapgen` when preparing for FFI.
     pub(crate) fn counter_regions(
         &self,
-    ) -> impl Iterator<Item = (Counter, &CodeRegion)> + ExactSizeIterator {
+    ) -> impl Iterator<Item = (MappingKind, &CodeRegion)> + ExactSizeIterator {
         self.function_coverage_info.mappings.iter().map(move |mapping| {
-            let &Mapping { term, ref code_region } = mapping;
-            let counter = self.counter_for_term(term);
-            (counter, code_region)
+            let Mapping { kind, code_region } = mapping;
+            let kind =
+                kind.map_terms(|term| if self.is_zero_term(term) { CovTerm::Zero } else { term });
+            (kind, code_region)
         })
     }
 
     fn counter_for_term(&self, term: CovTerm) -> Counter {
-        if is_zero_term(&self.counters_seen, &self.zero_expressions, term) {
-            Counter::ZERO
-        } else {
-            Counter::from_term(term)
-        }
+        if self.is_zero_term(term) { Counter::ZERO } else { Counter::from_term(term) }
+    }
+
+    fn is_zero_term(&self, term: CovTerm) -> bool {
+        is_zero_term(&self.counters_seen, &self.zero_expressions, term)
     }
 }
 

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -12,7 +12,6 @@ use rustc_hir::def_id::DefId;
 use rustc_index::IndexVec;
 use rustc_middle::bug;
 use rustc_middle::mir;
-use rustc_middle::mir::coverage::CodeRegion;
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_span::def_id::DefIdSet;
 use rustc_span::Symbol;
@@ -237,7 +236,7 @@ fn encode_mappings_for_function(
     // Prepare file IDs for each filename, and prepare the mapping data so that
     // we can pass it through FFI to LLVM.
     for (file_name, counter_regions_for_file) in
-        &counter_regions.group_by(|(_counter, region)| region.file_name)
+        &counter_regions.group_by(|(_, region)| region.file_name)
     {
         // Look up the global file ID for this filename.
         let global_file_id = global_file_table.global_file_id_for_file_name(file_name);
@@ -248,17 +247,12 @@ fn encode_mappings_for_function(
 
         // 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,
+        for (mapping_kind, region) in counter_regions_for_file {
+            debug!("Adding counter {mapping_kind:?} to map for {region:?}");
+            mapping_regions.push(CounterMappingRegion::from_mapping(
+                &mapping_kind,
                 local_file_id.as_u32(),
-                start_line,
-                start_col,
-                end_line,
-                end_col,
+                region,
             ));
         }
     }

compiler/rustc_middle/src/mir/coverage.rs

@@ -158,18 +158,36 @@ pub struct Expression {
     pub rhs: CovTerm,
 }
 
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+pub enum MappingKind {
+    /// Associates a normal region of code with a counter/expression/zero.
+    Code(CovTerm),
+}
+
+impl MappingKind {
+    /// Iterator over all coverage terms in this mapping kind.
+    pub fn terms(&self) -> impl Iterator<Item = CovTerm> {
+        let one = |a| std::iter::once(a);
+        match *self {
+            Self::Code(term) => one(term),
+        }
+    }
+
+    /// Returns a copy of this mapping kind, in which all coverage terms have
+    /// been replaced with ones returned by the given function.
+    pub fn map_terms(&self, map_fn: impl Fn(CovTerm) -> CovTerm) -> Self {
+        match *self {
+            Self::Code(term) => Self::Code(map_fn(term)),
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub struct Mapping {
+    pub kind: MappingKind,
     pub code_region: CodeRegion,
-
-    /// Indicates whether this mapping uses a counter value, expression value,
-    /// or zero value.
-    ///
-    /// FIXME: When we add support for mapping kinds other than `Code`
-    /// (e.g. branch regions, expansion regions), replace this with a dedicated
-    /// mapping-kind enum.
-    pub term: CovTerm,
 }
 
 /// Stores per-function coverage information attached to a `mir::Body`,

compiler/rustc_middle/src/mir/pretty.rs

@@ -493,8 +493,8 @@ fn write_function_coverage_info(
     for (id, expression) in expressions.iter_enumerated() {
         writeln!(w, "{INDENT}coverage {id:?} => {expression:?};")?;
     }
-    for coverage::Mapping { term, code_region } in mappings {
-        writeln!(w, "{INDENT}coverage {term:?} => {code_region:?};")?;
+    for coverage::Mapping { kind, code_region } in mappings {
+        writeln!(w, "{INDENT}coverage {kind:?} => {code_region:?};")?;
     }
     writeln!(w)?;
 

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -9,7 +9,7 @@ mod tests;
 
 use self::counters::{BcbCounter, CoverageCounters};
 use self::graph::{BasicCoverageBlock, CoverageGraph};
-use self::spans::CoverageSpans;
+use self::spans::{BcbMapping, BcbMappingKind, CoverageSpans};
 
 use crate::MirPass;
 
@@ -141,22 +141,21 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
         let file_name =
             Symbol::intern(&source_file.name.for_codegen(self.tcx.sess).to_string_lossy());
 
+        let term_for_bcb = |bcb| {
+            coverage_counters
+                .bcb_counter(bcb)
+                .expect("all BCBs with spans were given counters")
+                .as_term()
+        };
+
         coverage_spans
-            .bcbs_with_coverage_spans()
-            // For each BCB with spans, get a coverage term for its counter.
-            .map(|(bcb, spans)| {
-                let term = coverage_counters
-                    .bcb_counter(bcb)
-                    .expect("all BCBs with spans were given counters")
-                    .as_term();
-                (term, spans)
-            })
-            // Flatten the spans into individual term/span pairs.
-            .flat_map(|(term, spans)| spans.iter().map(move |&span| (term, span)))
-            // Convert each span to a code region, and create the final mapping.
-            .filter_map(|(term, span)| {
+            .all_bcb_mappings()
+            .filter_map(|&BcbMapping { kind: bcb_mapping_kind, span }| {
+                let kind = match bcb_mapping_kind {
+                    BcbMappingKind::Code(bcb) => MappingKind::Code(term_for_bcb(bcb)),
+                };
                 let code_region = make_code_region(source_map, file_name, span, body_span)?;
-                Some(Mapping { term, code_region })
+                Some(Mapping { kind, code_region })
             })
             .collect::<Vec<_>>()
     }

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -1,5 +1,5 @@
 use rustc_data_structures::graph::WithNumNodes;
-use rustc_index::IndexVec;
+use rustc_index::bit_set::BitSet;
 use rustc_middle::mir;
 use rustc_span::{BytePos, Span, DUMMY_SP};
 
@@ -8,9 +8,21 @@ use crate::coverage::ExtractedHirInfo;
 
 mod from_mir;
 
+#[derive(Clone, Copy, Debug)]
+pub(super) enum BcbMappingKind {
+    /// Associates an ordinary executable code span with its corresponding BCB.
+    Code(BasicCoverageBlock),
+}
+
+#[derive(Debug)]
+pub(super) struct BcbMapping {
+    pub(super) kind: BcbMappingKind,
+    pub(super) span: Span,
+}
+
 pub(super) struct CoverageSpans {
-    /// Map from BCBs to their list of coverage spans.
-    bcb_to_spans: IndexVec<BasicCoverageBlock, Vec<Span>>,
+    bcb_has_mappings: BitSet<BasicCoverageBlock>,
+    mappings: Vec<BcbMapping>,
 }
 
 impl CoverageSpans {
@@ -23,36 +35,42 @@ impl CoverageSpans {
         hir_info: &ExtractedHirInfo,
         basic_coverage_blocks: &CoverageGraph,
     ) -> Option<Self> {
+        let mut mappings = vec![];
+
         let coverage_spans = CoverageSpansGenerator::generate_coverage_spans(
             mir_body,
             hir_info,
             basic_coverage_blocks,
         );
+        mappings.extend(coverage_spans.into_iter().map(|CoverageSpan { bcb, span, .. }| {
+            // Each span produced by the generator represents an ordinary code region.
+            BcbMapping { kind: BcbMappingKind::Code(bcb), span }
+        }));
 
-        if coverage_spans.is_empty() {
+        if mappings.is_empty() {
             return None;
         }
 
-        // Group the coverage spans by BCB, with the BCBs in sorted order.
-        let mut bcb_to_spans = IndexVec::from_elem_n(Vec::new(), basic_coverage_blocks.num_nodes());
-        for CoverageSpan { bcb, span, .. } in coverage_spans {
-            bcb_to_spans[bcb].push(span);
+        // Identify which BCBs have one or more mappings.
+        let mut bcb_has_mappings = BitSet::new_empty(basic_coverage_blocks.num_nodes());
+        let mut insert = |bcb| {
+            bcb_has_mappings.insert(bcb);
+        };
+        for &BcbMapping { kind, span: _ } in &mappings {
+            match kind {
+                BcbMappingKind::Code(bcb) => insert(bcb),
+            }
         }
 
-        Some(Self { bcb_to_spans })
+        Some(Self { bcb_has_mappings, mappings })
     }
 
     pub(super) fn bcb_has_coverage_spans(&self, bcb: BasicCoverageBlock) -> bool {
-        !self.bcb_to_spans[bcb].is_empty()
+        self.bcb_has_mappings.contains(bcb)
     }
 
-    pub(super) fn bcbs_with_coverage_spans(
-        &self,
-    ) -> impl Iterator<Item = (BasicCoverageBlock, &[Span])> {
-        self.bcb_to_spans.iter_enumerated().filter_map(|(bcb, spans)| {
-            // Only yield BCBs that have at least one coverage span.
-            (!spans.is_empty()).then_some((bcb, spans.as_slice()))
-        })
+    pub(super) fn all_bcb_mappings(&self) -> impl Iterator<Item = &BcbMapping> {
+        self.mappings.iter()
     }
 }
 

tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff

@@ -4,7 +4,7 @@
   fn bar() -> bool {
       let mut _0: bool;
 
-+     coverage Counter(0) => /the/src/instrument_coverage.rs:21:1 - 23:2;
++     coverage Code(Counter(0)) => /the/src/instrument_coverage.rs:21:1 - 23:2;
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);

tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff

@@ -9,11 +9,11 @@
 
 +     coverage ExpressionId(0) => Expression { lhs: Counter(0), op: Add, rhs: Counter(1) };
 +     coverage ExpressionId(1) => Expression { lhs: Expression(0), op: Subtract, rhs: Counter(1) };
-+     coverage Counter(0) => /the/src/instrument_coverage.rs:12:1 - 12:11;
-+     coverage Expression(0) => /the/src/instrument_coverage.rs:13:5 - 14:17;
-+     coverage Expression(1) => /the/src/instrument_coverage.rs:15:13 - 15:18;
-+     coverage Expression(1) => /the/src/instrument_coverage.rs:18:1 - 18:2;
-+     coverage Counter(1) => /the/src/instrument_coverage.rs:16:10 - 16:11;
++     coverage Code(Counter(0)) => /the/src/instrument_coverage.rs:12:1 - 12:11;
++     coverage Code(Expression(0)) => /the/src/instrument_coverage.rs:13:5 - 14:17;
++     coverage Code(Expression(1)) => /the/src/instrument_coverage.rs:15:13 - 15:18;
++     coverage Code(Counter(1)) => /the/src/instrument_coverage.rs:16:10 - 16:11;
++     coverage Code(Expression(1)) => /the/src/instrument_coverage.rs:18:1 - 18:2;
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);