coverage: Restrict `ExpressionUsed` simplification to `Code` mappings

In the future, branch and MC/DC mappings might have expressions that don't
correspond to any single point in the control-flow graph. That makes it
trickier to keep track of which expressions should expect an `ExpressionUsed`
node.

We therefore sidestep that complexity by only performing `ExpressionUsed`
simplification for expressions associated directly with ordinary `Code`
mappings.

4 files changed, 27 insertions, 19 deletions

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
index b969fe27a99..14a94468587 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
@@ -66,8 +66,15 @@ 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 term in function_coverage_info.mappings.iter().flat_map(|m| m.kind.terms()) {
-            if let CovTerm::Expression(id) = term {
+        for mapping in function_coverage_info.mappings.iter() {
+            // Currently we only worry about ordinary code mappings.
+            // For branch and MC/DC mappings, expressions might not correspond
+            // to any particular point in the control-flow graph.
+            // (Keep this in sync with the injection of `ExpressionUsed`
+            // statements in the `InstrumentCoverage` MIR pass.)
+            if let MappingKind::Code(term) = mapping.kind
+                && let CovTerm::Expression(id) = term
+            {
                 expressions_seen.remove(id);
             }
         }
diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
index beaaadd497d..2a593340849 100644
--- a/compiler/rustc_middle/src/mir/coverage.rs
+++ b/compiler/rustc_middle/src/mir/coverage.rs
@@ -220,19 +220,6 @@ pub enum MappingKind {
 }
 
 impl MappingKind {
-    /// Iterator over all coverage terms in this mapping kind.
-    pub fn terms(&self) -> impl Iterator<Item = CovTerm> {
-        let zero = || None.into_iter().chain(None);
-        let one = |a| Some(a).into_iter().chain(None);
-        let two = |a, b| Some(a).into_iter().chain(Some(b));
-        match *self {
-            Self::Code(term) => one(term),
-            Self::Branch { true_term, false_term } => two(true_term, false_term),
-            Self::MCDCBranch { true_term, false_term, .. } => two(true_term, false_term),
-            Self::MCDCDecision(_) => zero(),
-        }
-    }
-
     /// 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 {
diff --git a/compiler/rustc_mir_transform/src/coverage/mappings.rs b/compiler/rustc_mir_transform/src/coverage/mappings.rs
index 578e1ae0b8a..2ac08ea85d2 100644
--- a/compiler/rustc_mir_transform/src/coverage/mappings.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mappings.rs
@@ -159,6 +159,15 @@ impl ExtractedMappings {
 
         bcbs_with_counter_mappings
     }
+
+    /// Returns the set of BCBs that have one or more `Code` mappings.
+    pub(super) fn bcbs_with_ordinary_code_mappings(&self) -> BitSet<BasicCoverageBlock> {
+        let mut bcbs = BitSet::new_empty(self.num_bcbs);
+        for &CodeMapping { span: _, bcb } in &self.code_mappings {
+            bcbs.insert(bcb);
+        }
+        bcbs
+    }
 }
 
 fn resolve_block_markers(
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 4e6692a55e5..3772a8f5118 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -25,7 +25,7 @@ use rustc_span::source_map::SourceMap;
 use rustc_span::{BytePos, Pos, RelativeBytePos, Span, Symbol};
 
 use crate::coverage::counters::{CounterIncrementSite, CoverageCounters};
-use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
+use crate::coverage::graph::CoverageGraph;
 use crate::coverage::mappings::ExtractedMappings;
 use crate::MirPass;
 
@@ -108,7 +108,7 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     inject_coverage_statements(
         mir_body,
         &basic_coverage_blocks,
-        bcb_has_counter_mappings,
+        &extracted_mappings,
         &coverage_counters,
     );
 
@@ -219,7 +219,7 @@ fn create_mappings<'tcx>(
 fn inject_coverage_statements<'tcx>(
     mir_body: &mut mir::Body<'tcx>,
     basic_coverage_blocks: &CoverageGraph,
-    bcb_has_coverage_spans: impl Fn(BasicCoverageBlock) -> bool,
+    extracted_mappings: &ExtractedMappings,
     coverage_counters: &CoverageCounters,
 ) {
     // Inject counter-increment statements into MIR.
@@ -252,11 +252,16 @@ fn inject_coverage_statements<'tcx>(
     // can check whether the injected statement survived MIR optimization.
     // (BCB edges can't have spans, so we only need to process BCB nodes here.)
     //
+    // We only do this for ordinary `Code` mappings, because branch and MC/DC
+    // mappings might have expressions that don't correspond to any single
+    // point in the control-flow graph.
+    //
     // See the code in `rustc_codegen_llvm::coverageinfo::map_data` that deals
     // with "expressions seen" and "zero terms".
+    let eligible_bcbs = extracted_mappings.bcbs_with_ordinary_code_mappings();
     for (bcb, expression_id) in coverage_counters
         .bcb_nodes_with_coverage_expressions()
-        .filter(|&(bcb, _)| bcb_has_coverage_spans(bcb))
+        .filter(|&(bcb, _)| eligible_bcbs.contains(bcb))
     {
         inject_statement(
             mir_body,