coverage: Store expression data in FunctionCoverageData

Even though expression details are now stored in the info structure, we still
need to inject `Expression` statements into MIR, because if one is missing
during codegen then we know that it was optimized out and we can remap all of
its associated code regions to zero.

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -2,19 +2,11 @@ use crate::coverageinfo::ffi::{Counter, CounterExpression, ExprKind};
 
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_index::bit_set::BitSet;
-use rustc_index::IndexVec;
 use rustc_middle::mir::coverage::{
-    CodeRegion, CounterId, CovTerm, ExpressionId, FunctionCoverageInfo, Mapping, Op,
+    CodeRegion, CounterId, CovTerm, Expression, ExpressionId, FunctionCoverageInfo, Mapping, Op,
 };
 use rustc_middle::ty::Instance;
 
-#[derive(Clone, Debug, PartialEq)]
-pub struct Expression {
-    lhs: CovTerm,
-    op: Op,
-    rhs: CovTerm,
-}
-
 /// Holds all of the coverage mapping data associated with a function instance,
 /// collected during traversal of `Coverage` statements in the function's MIR.
 #[derive(Debug)]
@@ -26,7 +18,9 @@ pub struct FunctionCoverage<'tcx> {
     /// Tracks which counters have been seen, so that we can identify mappings
     /// to counters that were optimized out, and set them to zero.
     counters_seen: BitSet<CounterId>,
-    expressions: IndexVec<ExpressionId, Option<Expression>>,
+    /// Tracks which expressions have one or more associated mappings, but have
+    /// not yet been seen. This lets us detect that they were optimized out.
+    expressions_not_seen: BitSet<ExpressionId>,
     /// Tracks which expressions are known to always have a value of zero.
     /// Only updated during the finalize step.
     zero_expressions: FxIndexSet<ExpressionId>,
@@ -55,16 +49,27 @@ impl<'tcx> FunctionCoverage<'tcx> {
         is_used: bool,
     ) -> Self {
         let num_counters = function_coverage_info.num_counters;
-        let num_expressions = function_coverage_info.num_expressions;
+        let num_expressions = function_coverage_info.expressions.len();
         debug!(
             "FunctionCoverage::create(instance={instance:?}) has \
             num_counters={num_counters}, num_expressions={num_expressions}, is_used={is_used}"
         );
+
+        // Every expression that has an associated mapping should have a
+        // corresponding statement in MIR. If we don't see that statement, we
+        // know that it was removed by MIR optimizations.
+        let mut expressions_not_seen = BitSet::new_empty(num_expressions);
+        for Mapping { term, .. } in &function_coverage_info.mappings {
+            if let &CovTerm::Expression(id) = term {
+                expressions_not_seen.insert(id);
+            }
+        }
+
         Self {
             function_coverage_info,
             is_used,
             counters_seen: BitSet::new_empty(num_counters),
-            expressions: IndexVec::from_elem_n(None, num_expressions),
+            expressions_not_seen,
             zero_expressions: FxIndexSet::default(),
         }
     }
@@ -80,35 +85,10 @@ impl<'tcx> FunctionCoverage<'tcx> {
         self.counters_seen.insert(id);
     }
 
-    /// Adds information about a coverage expression.
+    /// Marks an expression ID as having been seen.
     #[instrument(level = "debug", skip(self))]
-    pub(crate) fn add_counter_expression(
-        &mut self,
-        expression_id: ExpressionId,
-        lhs: CovTerm,
-        op: Op,
-        rhs: CovTerm,
-    ) {
-        debug_assert!(
-            expression_id.as_usize() < self.expressions.len(),
-            "expression_id {} is out of range for expressions.len() = {}
-            for {:?}",
-            expression_id.as_usize(),
-            self.expressions.len(),
-            self,
-        );
-
-        let expression = Expression { lhs, op, rhs };
-        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"
-            ),
-        }
+    pub(crate) fn add_counter_expression(&mut self, expression_id: ExpressionId) {
+        self.expressions_not_seen.remove(expression_id);
     }
 
     pub(crate) fn finalize(&mut self) {
@@ -133,13 +113,13 @@ impl<'tcx> FunctionCoverage<'tcx> {
         // and then update the set of always-zero expressions if necessary.
         // (By construction, expressions can only refer to other expressions
         // that have lower IDs, so one pass is sufficient.)
-        for (id, maybe_expression) in self.expressions.iter_enumerated() {
-            let Some(expression) = maybe_expression else {
-                // If an expression is missing, it must have been optimized away,
+        for (id, expression) in self.function_coverage_info.expressions.iter_enumerated() {
+            if self.expressions_not_seen.contains(id) {
+                // If an expression was not seen, it must have been optimized away,
                 // so any operand that refers to it can be replaced with zero.
                 zero_expressions.insert(id);
                 continue;
-            };
+            }
 
             // We don't need to simplify the actual expression data in the
             // expressions list; we can just simplify a temporary copy and then
@@ -191,17 +171,17 @@ impl<'tcx> FunctionCoverage<'tcx> {
             _ => Counter::from_term(operand),
         };
 
-        self.expressions
-            .iter()
-            .map(|expression| match expression {
-                None => {
+        self.function_coverage_info
+            .expressions
+            .iter_enumerated()
+            .map(|(id, &Expression { lhs, op, rhs })| {
+                if self.expressions_not_seen.contains(id) {
                     // This expression ID was allocated, but we never saw the
                     // actual expression, so it must have been optimized out.
                     // Replace it with a dummy expression, and let LLVM take
                     // care of omitting it from the expression list.
                     CounterExpression::DUMMY
-                }
-                &Some(Expression { lhs, op, rhs, .. }) => {
+                } else {
                     // Convert the operands and operator as normal.
                     CounterExpression::new(
                         counter_from_operand(lhs),

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -147,8 +147,8 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 );
                 bx.instrprof_increment(fn_name, hash, num_counters, index);
             }
-            CoverageKind::Expression { id, lhs, op, rhs } => {
-                func_coverage.add_counter_expression(id, lhs, op, rhs);
+            CoverageKind::Expression { id } => {
+                func_coverage.add_counter_expression(id);
             }
         }
     }

compiler/rustc_middle/src/mir/coverage.rs

@@ -1,5 +1,6 @@
 //! Metadata from source code coverage analysis and instrumentation.
 
+use rustc_index::IndexVec;
 use rustc_macros::HashStable;
 use rustc_span::Symbol;
 
@@ -70,9 +71,6 @@ pub enum CoverageKind {
         /// ID of this coverage-counter expression within its enclosing function.
         /// Other expressions in the same function can refer to it as an operand.
         id: ExpressionId,
-        lhs: CovTerm,
-        op: Op,
-        rhs: CovTerm,
     },
 }
 
@@ -81,17 +79,7 @@ impl Debug for CoverageKind {
         use CoverageKind::*;
         match self {
             Counter { id } => write!(fmt, "Counter({:?})", id.index()),
-            Expression { id, lhs, op, rhs } => write!(
-                fmt,
-                "Expression({:?}) = {:?} {} {:?}",
-                id.index(),
-                lhs,
-                match op {
-                    Op::Add => "+",
-                    Op::Subtract => "-",
-                },
-                rhs,
-            ),
+            Expression { id } => write!(fmt, "Expression({:?})", id.index()),
         }
     }
 }
@@ -166,7 +154,7 @@ pub struct Mapping {
 pub struct FunctionCoverageInfo {
     pub function_source_hash: u64,
     pub num_counters: usize,
-    pub num_expressions: usize,
 
+    pub expressions: IndexVec<ExpressionId, Expression>,
     pub mappings: Vec<Mapping>,
 }

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -142,10 +142,6 @@ impl CoverageCounters {
         self.next_counter_id.as_usize()
     }
 
-    pub(crate) fn num_expressions(&self) -> usize {
-        self.expressions.len()
-    }
-
     fn set_bcb_counter(
         &mut self,
         bcb: BasicCoverageBlock,

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -271,19 +271,10 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
         // files), which might be helpful in analyzing unused expressions, to still be generated.
         debug_used_expressions.alert_on_unused_expressions(&self.coverage_counters.debug_counters);
 
-        ////////////////////////////////////////////////////
-        // Finally, inject the intermediate expressions collected along the way.
-        for intermediate_expression in &self.coverage_counters.intermediate_expressions {
-            inject_intermediate_expression(
-                self.mir_body,
-                self.make_mir_coverage_kind(intermediate_expression),
-            );
-        }
-
         self.mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo {
             function_source_hash: self.function_source_hash,
             num_counters: self.coverage_counters.num_counters(),
-            num_expressions: self.coverage_counters.num_expressions(),
+            expressions: std::mem::take(&mut self.coverage_counters.expressions),
             mappings: std::mem::take(&mut self.mappings),
         }));
     }
@@ -411,10 +402,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
                         inject_to_bb,
                     );
                 }
-                BcbCounter::Expression { .. } => inject_intermediate_expression(
-                    self.mir_body,
-                    self.make_mir_coverage_kind(&counter_kind),
-                ),
+                BcbCounter::Expression { .. } => {}
             }
         }
     }
@@ -442,10 +430,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
     fn make_mir_coverage_kind(&self, counter_kind: &BcbCounter) -> CoverageKind {
         match *counter_kind {
             BcbCounter::Counter { id } => CoverageKind::Counter { id },
-            BcbCounter::Expression { id } => {
-                let Expression { lhs, op, rhs } = self.coverage_counters.expressions[id];
-                CoverageKind::Expression { id, lhs, op, rhs }
-            }
+            BcbCounter::Expression { id } => CoverageKind::Expression { id },
         }
     }
 }
@@ -484,20 +469,6 @@ fn inject_statement(mir_body: &mut mir::Body<'_>, counter_kind: CoverageKind, bb
     data.statements.insert(0, statement);
 }
 
-// Non-code expressions are injected into the coverage map, without generating executable code.
-fn inject_intermediate_expression(mir_body: &mut mir::Body<'_>, expression: CoverageKind) {
-    debug_assert!(matches!(expression, CoverageKind::Expression { .. }));
-    debug!("  injecting non-code expression {:?}", expression);
-    let inject_in_bb = mir::START_BLOCK;
-    let data = &mut mir_body[inject_in_bb];
-    let source_info = data.terminator().source_info;
-    let statement = Statement {
-        source_info,
-        kind: StatementKind::Coverage(Box::new(Coverage { kind: expression })),
-    };
-    data.statements.push(statement);
-}
-
 /// Convert the Span into its file name, start line and column, and end line and column
 fn make_code_region(
     source_map: &SourceMap,

tests/mir-opt/coverage_graphviz.main.InstrumentCoverage.0.dot

@@ -3,7 +3,7 @@ digraph Cov_0_3 {
     node [fontname="Courier, monospace"];
     edge [fontname="Courier, monospace"];
     bcb3__Cov_0_3 [shape="none", label=<<table border="0" cellborder="1" cellspacing="0"><tr><td bgcolor="gray" align="center" colspan="1">bcb3</td></tr><tr><td align="left" balign="left">Counter(bcb3) at 13:10-13:10<br align="left"/>    13:10-13:10: @5[0]: Coverage::Counter(1)</td></tr><tr><td align="left" balign="left">bb5: Goto</td></tr></table>>];
-    bcb2__Cov_0_3 [shape="none", label=<<table border="0" cellborder="1" cellspacing="0"><tr><td bgcolor="gray" align="center" colspan="1">bcb2</td></tr><tr><td align="left" balign="left">Expression(bcb1:(bcb0 + bcb3) - bcb3) at 12:13-12:18<br align="left"/>    12:13-12:18: @4[0]: Coverage::Expression(1) = Expression(0) - Counter(1)<br align="left"/>Expression(bcb1:(bcb0 + bcb3) - bcb3) at 15:2-15:2<br align="left"/>    15:2-15:2: @4.Return: return</td></tr><tr><td align="left" balign="left">bb4: Return</td></tr></table>>];
+    bcb2__Cov_0_3 [shape="none", label=<<table border="0" cellborder="1" cellspacing="0"><tr><td bgcolor="gray" align="center" colspan="1">bcb2</td></tr><tr><td align="left" balign="left">Expression(bcb1:(bcb0 + bcb3) - bcb3) at 12:13-12:18<br align="left"/>    12:13-12:18: @4[0]: Coverage::Expression(1)<br align="left"/>Expression(bcb1:(bcb0 + bcb3) - bcb3) at 15:2-15:2<br align="left"/>    15:2-15:2: @4.Return: return</td></tr><tr><td align="left" balign="left">bb4: Return</td></tr></table>>];
     bcb1__Cov_0_3 [shape="none", label=<<table border="0" cellborder="1" cellspacing="0"><tr><td bgcolor="gray" align="center" colspan="1">bcb1</td></tr><tr><td align="left" balign="left">Expression(bcb0 + bcb3) at 10:5-11:17<br align="left"/>    11:12-11:17: @2.Call: _2 = bar() -&gt; [return: bb3, unwind: bb6]</td></tr><tr><td align="left" balign="left">bb1: FalseUnwind<br align="left"/>bb2: Call</td></tr><tr><td align="left" balign="left">bb3: SwitchInt</td></tr></table>>];
     bcb0__Cov_0_3 [shape="none", label=<<table border="0" cellborder="1" cellspacing="0"><tr><td bgcolor="gray" align="center" colspan="1">bcb0</td></tr><tr><td align="left" balign="left"></td></tr><tr><td align="left" balign="left">Counter(bcb0) at 9:1-9:11<br align="left"/>    </td></tr><tr><td align="left" balign="left">bb0: Goto</td></tr></table>>];
     bcb3__Cov_0_3 -> bcb1__Cov_0_3 [label=<>];

tests/mir-opt/instrument_coverage.main.InstrumentCoverage.diff

@@ -13,7 +13,7 @@
       }
 
       bb1: {
-+         Coverage::Expression(0) = Counter(0) + Counter(1);
++         Coverage::Expression(0);
           falseUnwind -> [real: bb2, unwind: bb6];
       }
 
@@ -27,7 +27,7 @@
       }
 
       bb4: {
-+         Coverage::Expression(1) = Expression(0) - Counter(1);
++         Coverage::Expression(1);
           _0 = const ();
           StorageDead(_2);
           return;