3 files changed, 39 insertions, 25 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,
             ));
         }
     }