coverage: Add enums to accommodate other kinds of coverage mappings

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);