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

9 files changed, 122 insertions, 73 deletions

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs
index 7ad2d03a5ed..017843c7e7d 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs
+++ b/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,
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
index cd67fafb8e4..d85d9411f03 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
+++ b/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)
     }
 }
 
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
index 51df14df644..6116a6fd222 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
+++ b/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,
             ));
         }
     }
diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
index ec5edceb269..18e198eb9fc 100644
--- a/compiler/rustc_middle/src/mir/coverage.rs
+++ b/compiler/rustc_middle/src/mir/coverage.rs
@@ -160,16 +160,34 @@ pub struct Expression {
 
 #[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`,
diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
index a1e5d73a0fd..1a6b0f4031d 100644
--- a/compiler/rustc_middle/src/mir/pretty.rs
+++ b/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)?;
 
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index dcd7014f4fc..a11d224e8f1 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/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<_>>()
     }
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index 5983189984d..81f6c831206 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/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()
     }
 }
 
diff --git a/tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff b/tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff
index 1ef6b69ef5b..dec99ff4601 100644
--- a/tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff
+++ b/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);
diff --git a/tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff b/tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff
index 14b4833a515..368088d1296 100644
--- a/tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff
+++ b/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);