coverage: Let each coverage statement hold a vector of code regions

This makes it possible for a `StatementKind::Coverage` to hold more than one
code region, but that capability is not yet used.

9 files changed, 96 insertions, 90 deletions

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
index e83110dcad4..55f43aa5341 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
@@ -11,7 +11,7 @@ pub struct Expression {
     lhs: Operand,
     op: Op,
     rhs: Operand,
-    region: Option<CodeRegion>,
+    code_regions: Vec<CodeRegion>,
 }
 
 /// Collects all of the coverage regions associated with (a) injected counters, (b) counter
@@ -30,7 +30,7 @@ pub struct FunctionCoverage<'tcx> {
     instance: Instance<'tcx>,
     source_hash: u64,
     is_used: bool,
-    counters: IndexVec<CounterId, Option<CodeRegion>>,
+    counters: IndexVec<CounterId, Option<Vec<CodeRegion>>>,
     expressions: IndexVec<ExpressionId, Option<Expression>>,
     unreachable_regions: Vec<CodeRegion>,
 }
@@ -77,28 +77,40 @@ impl<'tcx> FunctionCoverage<'tcx> {
         }
     }
 
-    /// Adds a code region to be counted by an injected counter intrinsic.
-    pub fn add_counter(&mut self, id: CounterId, region: CodeRegion) {
-        if let Some(previous_region) = self.counters[id].replace(region.clone()) {
-            assert_eq!(previous_region, region, "add_counter: code region for id changed");
+    /// Adds code regions to be counted by an injected counter intrinsic.
+    #[instrument(level = "debug", skip(self))]
+    pub(crate) fn add_counter(&mut self, id: CounterId, code_regions: &[CodeRegion]) {
+        if code_regions.is_empty() {
+            return;
+        }
+
+        let slot = &mut self.counters[id];
+        match slot {
+            None => *slot = Some(code_regions.to_owned()),
+            // If this counter ID slot has already been filled, it should
+            // contain identical information.
+            Some(ref previous_regions) => assert_eq!(
+                previous_regions, code_regions,
+                "add_counter: code regions for id changed"
+            ),
         }
     }
 
+    /// Adds information about a coverage expression, along with zero or more
+    /// code regions mapped to that expression.
+    ///
     /// Both counters and "counter expressions" (or simply, "expressions") can be operands in other
     /// expressions. These are tracked as separate variants of `Operand`, so there is no ambiguity
     /// between operands that are counter IDs and operands that are expression IDs.
-    pub fn add_counter_expression(
+    #[instrument(level = "debug", skip(self))]
+    pub(crate) fn add_counter_expression(
         &mut self,
         expression_id: ExpressionId,
         lhs: Operand,
         op: Op,
         rhs: Operand,
-        region: Option<CodeRegion>,
+        code_regions: &[CodeRegion],
     ) {
-        debug!(
-            "add_counter_expression({:?}, lhs={:?}, op={:?}, rhs={:?} at {:?}",
-            expression_id, lhs, op, rhs, region
-        );
         debug_assert!(
             expression_id.as_usize() < self.expressions.len(),
             "expression_id {} is out of range for expressions.len() = {}
@@ -107,23 +119,25 @@ impl<'tcx> FunctionCoverage<'tcx> {
             self.expressions.len(),
             self,
         );
-        if let Some(previous_expression) = self.expressions[expression_id].replace(Expression {
-            lhs,
-            op,
-            rhs,
-            region: region.clone(),
-        }) {
-            assert_eq!(
-                previous_expression,
-                Expression { lhs, op, rhs, region },
+
+        let expression = Expression { lhs, op, rhs, code_regions: code_regions.to_owned() };
+        let slot = &mut self.expressions[expression_id];
+        match slot {
+            None => *slot = Some(expression),
+            // If this expression ID slot has already been filled, it should
+            // contain identical information.
+            Some(ref previous_expression) => assert_eq!(
+                previous_expression, &expression,
                 "add_counter_expression: expression for id changed"
-            );
+            ),
         }
     }
 
-    /// Add a region that will be marked as "unreachable", with a constant "zero counter".
-    pub fn add_unreachable_region(&mut self, region: CodeRegion) {
-        self.unreachable_regions.push(region)
+    /// Adds regions that will be marked as "unreachable", with a constant "zero counter".
+    #[instrument(level = "debug", skip(self))]
+    pub(crate) fn add_unreachable_regions(&mut self, code_regions: &[CodeRegion]) {
+        assert!(!code_regions.is_empty(), "unreachable regions always have code regions");
+        self.unreachable_regions.extend_from_slice(code_regions);
     }
 
     /// Perform some simplifications to make the final coverage mappings
@@ -212,11 +226,16 @@ impl<'tcx> FunctionCoverage<'tcx> {
     }
 
     fn counter_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.counters.iter_enumerated().filter_map(|(index, entry)| {
-            // Option::map() will return None to filter out missing counters. This may happen
-            // if, for example, a MIR-instrumented counter is removed during an optimization.
-            entry.as_ref().map(|region| (Counter::counter_value_reference(index), region))
-        })
+        self.counters
+            .iter_enumerated()
+            // Filter out counter IDs that we never saw during MIR traversal.
+            // This can happen if a counter was optimized out by MIR transforms
+            // (and replaced with `CoverageKind::Unreachable` instead).
+            .filter_map(|(id, maybe_code_regions)| Some((id, maybe_code_regions.as_ref()?)))
+            .flat_map(|(id, code_regions)| {
+                let counter = Counter::counter_value_reference(id);
+                code_regions.iter().map(move |region| (counter, region))
+            })
     }
 
     /// Convert this function's coverage expression data into a form that can be
@@ -254,13 +273,17 @@ impl<'tcx> FunctionCoverage<'tcx> {
 
     fn expression_regions(&self) -> Vec<(Counter, &CodeRegion)> {
         // Find all of the expression IDs that weren't optimized out AND have
-        // an attached code region, and return the corresponding mapping as a
-        // counter/region pair.
+        // one or more attached code regions, and return the corresponding
+        // mappings as counter/region pairs.
         self.expressions
             .iter_enumerated()
-            .filter_map(|(id, expression)| {
-                let code_region = expression.as_ref()?.region.as_ref()?;
-                Some((Counter::expression(id), code_region))
+            .filter_map(|(id, maybe_expression)| {
+                let code_regions = &maybe_expression.as_ref()?.code_regions;
+                Some((id, code_regions))
+            })
+            .flat_map(|(id, code_regions)| {
+                let counter = Counter::expression(id);
+                code_regions.iter().map(move |code_region| (counter, code_region))
             })
             .collect::<Vec<_>>()
     }
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
index f0ccaf4a3cd..dd2ce9b525b 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
@@ -108,25 +108,15 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
             .entry(instance)
             .or_insert_with(|| FunctionCoverage::new(bx.tcx(), instance));
 
-        let Coverage { kind, code_region } = coverage.clone();
-        match kind {
+        let Coverage { kind, code_regions } = coverage;
+        match *kind {
             CoverageKind::Counter { function_source_hash, id } => {
                 debug!(
                     "ensuring function source hash is set for instance={:?}; function_source_hash={}",
                     instance, function_source_hash,
                 );
                 func_coverage.set_function_source_hash(function_source_hash);
-
-                if let Some(code_region) = code_region {
-                    // Note: Some counters do not have code regions, but may still be referenced
-                    // from expressions. In that case, don't add the counter to the coverage map,
-                    // but do inject the counter intrinsic.
-                    debug!(
-                        "adding counter to coverage_map: instance={:?}, id={:?}, region={:?}",
-                        instance, id, code_region,
-                    );
-                    func_coverage.add_counter(id, code_region);
-                }
+                func_coverage.add_counter(id, code_regions);
                 // We need to explicitly drop the `RefMut` before calling into `instrprof_increment`,
                 // as that needs an exclusive borrow.
                 drop(coverage_map);
@@ -144,20 +134,10 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 bx.instrprof_increment(fn_name, hash, num_counters, index);
             }
             CoverageKind::Expression { id, lhs, op, rhs } => {
-                debug!(
-                    "adding counter expression to coverage_map: instance={:?}, id={:?}, {:?} {:?} {:?}; region: {:?}",
-                    instance, id, lhs, op, rhs, code_region,
-                );
-                func_coverage.add_counter_expression(id, lhs, op, rhs, code_region);
+                func_coverage.add_counter_expression(id, lhs, op, rhs, code_regions);
             }
             CoverageKind::Unreachable => {
-                let code_region =
-                    code_region.expect("unreachable regions always have code regions");
-                debug!(
-                    "adding unreachable code to coverage_map: instance={:?}, at {:?}",
-                    instance, code_region,
-                );
-                func_coverage.add_unreachable_region(code_region);
+                func_coverage.add_unreachable_regions(code_regions);
             }
         }
     }
@@ -226,7 +206,8 @@ fn add_unused_function_coverage<'tcx>(
 
     let mut function_coverage = FunctionCoverage::unused(tcx, instance);
     for &code_region in tcx.covered_code_regions(def_id) {
-        function_coverage.add_unreachable_region(code_region.clone());
+        let code_region = std::slice::from_ref(code_region);
+        function_coverage.add_unreachable_regions(code_region);
     }
 
     if let Some(coverage_context) = cx.coverage_context() {
diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
index 76567c3f6b0..00998386b6b 100644
--- a/compiler/rustc_middle/src/mir/pretty.rs
+++ b/compiler/rustc_middle/src/mir/pretty.rs
@@ -16,7 +16,7 @@ use rustc_middle::mir::interpret::{
     Pointer, Provenance,
 };
 use rustc_middle::mir::visit::Visitor;
-use rustc_middle::mir::*;
+use rustc_middle::mir::{self, *};
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_target::abi::Size;
 
@@ -685,10 +685,13 @@ impl Debug for Statement<'_> {
             AscribeUserType(box (ref place, ref c_ty), ref variance) => {
                 write!(fmt, "AscribeUserType({place:?}, {variance:?}, {c_ty:?})")
             }
-            Coverage(box self::Coverage { ref kind, code_region: Some(ref rgn) }) => {
-                write!(fmt, "Coverage::{kind:?} for {rgn:?}")
+            Coverage(box mir::Coverage { ref kind, ref code_regions }) => {
+                if code_regions.is_empty() {
+                    write!(fmt, "Coverage::{kind:?}")
+                } else {
+                    write!(fmt, "Coverage::{kind:?} for {code_regions:?}")
+                }
             }
-            Coverage(box ref coverage) => write!(fmt, "Coverage::{:?}", coverage.kind),
             Intrinsic(box ref intrinsic) => write!(fmt, "{intrinsic}"),
             ConstEvalCounter => write!(fmt, "ConstEvalCounter"),
             Nop => write!(fmt, "nop"),
diff --git a/compiler/rustc_middle/src/mir/syntax.rs b/compiler/rustc_middle/src/mir/syntax.rs
index 55f895f73b4..1bd73f9e34d 100644
--- a/compiler/rustc_middle/src/mir/syntax.rs
+++ b/compiler/rustc_middle/src/mir/syntax.rs
@@ -514,7 +514,7 @@ pub enum FakeReadCause {
 #[derive(TypeFoldable, TypeVisitable)]
 pub struct Coverage {
     pub kind: CoverageKind,
-    pub code_region: Option<CodeRegion>,
+    pub code_regions: Vec<CodeRegion>,
 }
 
 #[derive(Clone, Debug, PartialEq, TyEncodable, TyDecodable, Hash, HashStable)]
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 6d7f8b5302b..967b5766e95 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -243,7 +243,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
                 self.mir_body,
                 self.make_mir_coverage_kind(&counter_kind),
                 self.bcb_leader_bb(bcb),
-                Some(code_region),
+                vec![code_region],
             );
         }
     }
@@ -302,7 +302,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
                         self.mir_body,
                         self.make_mir_coverage_kind(&counter_kind),
                         inject_to_bb,
-                        None,
+                        Vec::new(),
                     );
                 }
                 BcbCounter::Expression { .. } => inject_intermediate_expression(
@@ -367,20 +367,14 @@ fn inject_statement(
     mir_body: &mut mir::Body<'_>,
     counter_kind: CoverageKind,
     bb: BasicBlock,
-    some_code_region: Option<CodeRegion>,
+    code_regions: Vec<CodeRegion>,
 ) {
-    debug!(
-        "  injecting statement {:?} for {:?} at code region: {:?}",
-        counter_kind, bb, some_code_region
-    );
+    debug!("  injecting statement {counter_kind:?} for {bb:?} at code regions: {code_regions:?}");
     let data = &mut mir_body[bb];
     let source_info = data.terminator().source_info;
     let statement = Statement {
         source_info,
-        kind: StatementKind::Coverage(Box::new(Coverage {
-            kind: counter_kind,
-            code_region: some_code_region,
-        })),
+        kind: StatementKind::Coverage(Box::new(Coverage { kind: counter_kind, code_regions })),
     };
     data.statements.insert(0, statement);
 }
@@ -394,7 +388,10 @@ fn inject_intermediate_expression(mir_body: &mut mir::Body<'_>, expression: Cove
     let source_info = data.terminator().source_info;
     let statement = Statement {
         source_info,
-        kind: StatementKind::Coverage(Box::new(Coverage { kind: expression, code_region: None })),
+        kind: StatementKind::Coverage(Box::new(Coverage {
+            kind: expression,
+            code_regions: Vec::new(),
+        })),
     };
     data.statements.push(statement);
 }
diff --git a/compiler/rustc_mir_transform/src/coverage/query.rs b/compiler/rustc_mir_transform/src/coverage/query.rs
index 56365c5d474..2c0164e765c 100644
--- a/compiler/rustc_mir_transform/src/coverage/query.rs
+++ b/compiler/rustc_mir_transform/src/coverage/query.rs
@@ -93,8 +93,8 @@ fn coverageinfo<'tcx>(tcx: TyCtxt<'tcx>, instance_def: ty::InstanceDef<'tcx>) ->
 fn covered_code_regions(tcx: TyCtxt<'_>, def_id: DefId) -> Vec<&CodeRegion> {
     let body = mir_body(tcx, def_id);
     all_coverage_in_mir_body(body)
-        // Not all coverage statements have an attached code region.
-        .filter_map(|coverage| coverage.code_region.as_ref())
+        // Coverage statements have a list of code regions (possibly empty).
+        .flat_map(|coverage| coverage.code_regions.as_slice())
         .collect()
 }
 
diff --git a/compiler/rustc_mir_transform/src/simplify.rs b/compiler/rustc_mir_transform/src/simplify.rs
index 2795cf15702..73dae044355 100644
--- a/compiler/rustc_mir_transform/src/simplify.rs
+++ b/compiler/rustc_mir_transform/src/simplify.rs
@@ -441,21 +441,23 @@ fn save_unreachable_coverage(
         let dead_block = &basic_blocks[dead_block];
         for statement in &dead_block.statements {
             let StatementKind::Coverage(coverage) = &statement.kind else { continue };
-            let Some(code_region) = &coverage.code_region else { continue };
+            if coverage.code_regions.is_empty() {
+                continue;
+            };
             let instance = statement.source_info.scope.inlined_instance(source_scopes);
             if live.contains(&instance) {
-                retained_coverage.push((statement.source_info, code_region.clone()));
+                retained_coverage.push((statement.source_info, coverage.code_regions.clone()));
             }
         }
     }
 
     let start_block = &mut basic_blocks[START_BLOCK];
     start_block.statements.extend(retained_coverage.into_iter().map(
-        |(source_info, code_region)| Statement {
+        |(source_info, code_regions)| Statement {
             source_info,
             kind: StatementKind::Coverage(Box::new(Coverage {
                 kind: CoverageKind::Unreachable,
-                code_region: Some(code_region),
+                code_regions,
             })),
         },
     ));
diff --git a/tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff b/tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff
index afcfde09c02..28a7ffda371 100644
--- a/tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff
+++ b/tests/mir-opt/instrument_coverage.bar.InstrumentCoverage.diff
@@ -5,7 +5,7 @@
       let mut _0: bool;
   
       bb0: {
-+         Coverage::Counter(0) for /the/src/instrument_coverage.rs:20:1 - 22:2;
++         Coverage::Counter(0) for [/the/src/instrument_coverage.rs:20:1 - 22:2];
           _0 = const true;
           return;
       }
diff --git a/tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff b/tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff
index e17c6ddc56e..756e5a59b3f 100644
--- a/tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff
+++ b/tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff
@@ -8,12 +8,12 @@
       let mut _3: !;
   
       bb0: {
-+         Coverage::Counter(0) for /the/src/instrument_coverage.rs:11:1 - 11:11;
++         Coverage::Counter(0) for [/the/src/instrument_coverage.rs:11:1 - 11:11];
           goto -> bb1;
       }
   
       bb1: {
-+         Coverage::Expression(0) = Counter(0) + Counter(1) for /the/src/instrument_coverage.rs:12:5 - 13:17;
++         Coverage::Expression(0) = Counter(0) + Counter(1) for [/the/src/instrument_coverage.rs:12:5 - 13:17];
           falseUnwind -> [real: bb2, unwind: bb6];
       }
   
@@ -27,15 +27,15 @@
       }
   
       bb4: {
-+         Coverage::Expression(2) = Expression(1) + Zero for /the/src/instrument_coverage.rs:17:1 - 17:2;
-+         Coverage::Expression(1) = Expression(0) - Counter(1) for /the/src/instrument_coverage.rs:14:13 - 14:18;
++         Coverage::Expression(2) = Expression(1) + Zero for [/the/src/instrument_coverage.rs:17:1 - 17:2];
++         Coverage::Expression(1) = Expression(0) - Counter(1) for [/the/src/instrument_coverage.rs:14:13 - 14:18];
           _0 = const ();
           StorageDead(_2);
           return;
       }
   
       bb5: {
-+         Coverage::Counter(1) for /the/src/instrument_coverage.rs:15:10 - 15:11;
++         Coverage::Counter(1) for [/the/src/instrument_coverage.rs:15:10 - 15:11];
           _1 = const ();
           StorageDead(_2);
           goto -> bb1;