1 files changed, 102 insertions, 5 deletions
diff --git a/compiler/rustc_mir_transform/src/coverage/query.rs b/compiler/rustc_mir_transform/src/coverage/query.rs
index 0090f6f3040..edaec3c7965 100644
--- a/compiler/rustc_mir_transform/src/coverage/query.rs
+++ b/compiler/rustc_mir_transform/src/coverage/query.rs
@@ -1,8 +1,11 @@
 use rustc_data_structures::captures::Captures;
 use rustc_index::bit_set::BitSet;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
-use rustc_middle::mir::coverage::{CovTerm, CoverageKind, MappingKind};
-use rustc_middle::mir::{Body, CoverageIdsInfo, Statement, StatementKind};
+use rustc_middle::mir::coverage::{
+    CounterId, CovTerm, CoverageIdsInfo, CoverageKind, Expression, ExpressionId,
+    FunctionCoverageInfo, MappingKind, Op,
+};
+use rustc_middle::mir::{Body, Statement, StatementKind};
 use rustc_middle::query::TyCtxtAt;
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_middle::util::Providers;
@@ -87,10 +90,10 @@ fn coverage_ids_info<'tcx>(
 ) -> CoverageIdsInfo {
     let mir_body = tcx.instance_mir(instance_def);
 
-    let Some(fn_cov_info) = mir_body.function_coverage_info.as_ref() else {
+    let Some(fn_cov_info) = mir_body.function_coverage_info.as_deref() else {
         return CoverageIdsInfo {
             counters_seen: BitSet::new_empty(0),
-            expressions_seen: BitSet::new_empty(0),
+            zero_expressions: BitSet::new_empty(0),
         };
     };
 
@@ -123,7 +126,10 @@ fn coverage_ids_info<'tcx>(
         }
     }
 
-    CoverageIdsInfo { counters_seen, expressions_seen }
+    let zero_expressions =
+        identify_zero_expressions(fn_cov_info, &counters_seen, &expressions_seen);
+
+    CoverageIdsInfo { counters_seen, zero_expressions }
 }
 
 fn all_coverage_in_mir_body<'a, 'tcx>(
@@ -141,3 +147,94 @@ fn is_inlined(body: &Body<'_>, statement: &Statement<'_>) -> bool {
     let scope_data = &body.source_scopes[statement.source_info.scope];
     scope_data.inlined.is_some() || scope_data.inlined_parent_scope.is_some()
 }
+
+/// Identify expressions that will always have a value of zero, and note
+/// their IDs in a `BitSet`. Mappings that refer to a zero expression
+/// can instead become mappings to a constant zero value.
+///
+/// This function mainly exists to preserve the simplifications that were
+/// already being performed by the Rust-side expression renumbering, so that
+/// the resulting coverage mappings don't get worse.
+fn identify_zero_expressions(
+    fn_cov_info: &FunctionCoverageInfo,
+    counters_seen: &BitSet<CounterId>,
+    expressions_seen: &BitSet<ExpressionId>,
+) -> BitSet<ExpressionId> {
+    // The set of expressions that either were optimized out entirely, or
+    // have zero as both of their operands, and will therefore always have
+    // a value of zero. Other expressions that refer to these as operands
+    // can have those operands replaced with `CovTerm::Zero`.
+    let mut zero_expressions = BitSet::new_empty(fn_cov_info.expressions.len());
+
+    // Simplify a copy of each expression based on lower-numbered expressions,
+    // 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, expression) in fn_cov_info.expressions.iter_enumerated() {
+        if !expressions_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
+        // use that to update the set of always-zero expressions.
+        let Expression { mut lhs, op, mut rhs } = *expression;
+
+        // If an expression has an operand that is also an expression, the
+        // operand's ID must be strictly lower. This is what lets us find
+        // all zero expressions in one pass.
+        let assert_operand_expression_is_lower = |operand_id: ExpressionId| {
+            assert!(
+                operand_id < id,
+                "Operand {operand_id:?} should be less than {id:?} in {expression:?}",
+            )
+        };
+
+        // If an operand refers to a counter or expression that is always
+        // zero, then that operand can be replaced with `CovTerm::Zero`.
+        let maybe_set_operand_to_zero = |operand: &mut CovTerm| {
+            if let CovTerm::Expression(id) = *operand {
+                assert_operand_expression_is_lower(id);
+            }
+
+            if is_zero_term(&counters_seen, &zero_expressions, *operand) {
+                *operand = CovTerm::Zero;
+            }
+        };
+        maybe_set_operand_to_zero(&mut lhs);
+        maybe_set_operand_to_zero(&mut rhs);
+
+        // Coverage counter values cannot be negative, so if an expression
+        // involves subtraction from zero, assume that its RHS must also be zero.
+        // (Do this after simplifications that could set the LHS to zero.)
+        if lhs == CovTerm::Zero && op == Op::Subtract {
+            rhs = CovTerm::Zero;
+        }
+
+        // After the above simplifications, if both operands are zero, then
+        // we know that this expression is always zero too.
+        if lhs == CovTerm::Zero && rhs == CovTerm::Zero {
+            zero_expressions.insert(id);
+        }
+    }
+
+    zero_expressions
+}
+
+/// Returns `true` if the given term is known to have a value of zero, taking
+/// into account knowledge of which counters are unused and which expressions
+/// are always zero.
+fn is_zero_term(
+    counters_seen: &BitSet<CounterId>,
+    zero_expressions: &BitSet<ExpressionId>,
+    term: CovTerm,
+) -> bool {
+    match term {
+        CovTerm::Zero => true,
+        CovTerm::Counter(id) => !counters_seen.contains(id),
+        CovTerm::Expression(id) => zero_expressions.contains(id),
+    }
+}