coverage: Store expression data in function coverage info

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>,
 }
 
 impl<'tcx> FunctionCoverage<'tcx> {
@@ -52,16 +46,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,
         }
     }
 
@@ -76,35 +81,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);
     }
 
     /// Identify expressions that will always have a value of zero, and note
@@ -125,13 +105,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
@@ -184,7 +164,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
         // Expression IDs are indices into `self.expressions`, and on the LLVM
         // side they will be treated as indices into `counter_expressions`, so
         // the two vectors should correspond 1:1.
-        assert_eq!(self.expressions.len(), counter_expressions.len());
+        assert_eq!(self.function_coverage_info.expressions.len(), counter_expressions.len());
 
         let counter_regions = self.counter_regions(zero_expressions);
 
@@ -204,17 +184,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()),
         }
     }
 }
@@ -133,6 +121,14 @@ impl Op {
     }
 }
 
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+pub struct Expression {
+    pub lhs: CovTerm,
+    pub op: Op,
+    pub rhs: CovTerm,
+}
+
 #[derive(Clone, Debug)]
 #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub struct Mapping {
@@ -155,7 +151,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

@@ -19,7 +19,7 @@ const NESTED_INDENT: &str = "    ";
 #[derive(Clone)]
 pub(super) enum BcbCounter {
     Counter { id: CounterId },
-    Expression { id: ExpressionId, lhs: CovTerm, op: Op, rhs: CovTerm },
+    Expression { id: ExpressionId },
 }
 
 impl BcbCounter {
@@ -39,17 +39,7 @@ impl Debug for BcbCounter {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Self::Counter { id, .. } => write!(fmt, "Counter({:?})", id.index()),
-            Self::Expression { id, lhs, op, rhs } => write!(
-                fmt,
-                "Expression({:?}) = {:?} {} {:?}",
-                id.index(),
-                lhs,
-                match op {
-                    Op::Add => "+",
-                    Op::Subtract => "-",
-                },
-                rhs,
-            ),
+            Self::Expression { id } => write!(fmt, "Expression({:?})", id.index()),
         }
     }
 }
@@ -58,7 +48,6 @@ impl Debug for BcbCounter {
 /// associated with nodes/edges in the BCB graph.
 pub(super) struct CoverageCounters {
     next_counter_id: CounterId,
-    next_expression_id: ExpressionId,
 
     /// Coverage counters/expressions that are associated with individual BCBs.
     bcb_counters: IndexVec<BasicCoverageBlock, Option<BcbCounter>>,
@@ -69,10 +58,9 @@ pub(super) struct CoverageCounters {
     /// Only used by debug assertions, to verify that BCBs with incoming edge
     /// counters do not have their own physical counters (expressions are allowed).
     bcb_has_incoming_edge_counters: BitSet<BasicCoverageBlock>,
-    /// Expression nodes that are not directly associated with any particular
-    /// BCB/edge, but are needed as operands to more complex expressions.
-    /// These are always [`BcbCounter::Expression`].
-    pub(super) intermediate_expressions: Vec<BcbCounter>,
+    /// Table of expression data, associating each expression ID with its
+    /// corresponding operator (+ or -) and its LHS/RHS operands.
+    pub(super) expressions: IndexVec<ExpressionId, Expression>,
 }
 
 impl CoverageCounters {
@@ -81,12 +69,10 @@ impl CoverageCounters {
 
         Self {
             next_counter_id: CounterId::START,
-            next_expression_id: ExpressionId::START,
-
             bcb_counters: IndexVec::from_elem_n(None, num_bcbs),
             bcb_edge_counters: FxHashMap::default(),
             bcb_has_incoming_edge_counters: BitSet::new_empty(num_bcbs),
-            intermediate_expressions: Vec::new(),
+            expressions: IndexVec::new(),
         }
     }
 
@@ -107,8 +93,8 @@ impl CoverageCounters {
     }
 
     fn make_expression(&mut self, lhs: CovTerm, op: Op, rhs: CovTerm) -> BcbCounter {
-        let id = self.next_expression();
-        BcbCounter::Expression { id, lhs, op, rhs }
+        let id = self.expressions.push(Expression { lhs, op, rhs });
+        BcbCounter::Expression { id }
     }
 
     /// Counter IDs start from one and go up.
@@ -118,22 +104,10 @@ impl CoverageCounters {
         next
     }
 
-    /// Expression IDs start from 0 and go up.
-    /// (Counter IDs and Expression IDs are distinguished by the `Operand` enum.)
-    fn next_expression(&mut self) -> ExpressionId {
-        let next = self.next_expression_id;
-        self.next_expression_id = self.next_expression_id + 1;
-        next
-    }
-
     pub(super) fn num_counters(&self) -> usize {
         self.next_counter_id.as_usize()
     }
 
-    pub(super) fn num_expressions(&self) -> usize {
-        self.next_expression_id.as_usize()
-    }
-
     fn set_bcb_counter(
         &mut self,
         bcb: BasicCoverageBlock,
@@ -333,7 +307,6 @@ impl<'a> MakeBcbCounters<'a> {
                     );
                     debug!("  [new intermediate expression: {:?}]", intermediate_expression);
                     let intermediate_expression_operand = intermediate_expression.as_term();
-                    self.coverage_counters.intermediate_expressions.push(intermediate_expression);
                     some_sumup_counter_operand.replace(intermediate_expression_operand);
                 }
             }
@@ -446,7 +419,6 @@ impl<'a> MakeBcbCounters<'a> {
                     intermediate_expression
                 );
                 let intermediate_expression_operand = intermediate_expression.as_term();
-                self.coverage_counters.intermediate_expressions.push(intermediate_expression);
                 some_sumup_edge_counter_operand.replace(intermediate_expression_operand);
             }
         }