Generalize Coverage Expression simplification

This extends the current simplification code to not only replace operands by `Zero`, but also to remove trivial `Counter + Zero` expressions and replace those with just `Counter`.

Author: Swatinem

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -1,6 +1,6 @@
 use crate::coverageinfo::ffi::{Counter, CounterExpression, ExprKind};
 
-use rustc_data_structures::fx::FxIndexSet;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_index::IndexVec;
 use rustc_middle::mir::coverage::{CodeRegion, CounterId, ExpressionId, Op, Operand};
 use rustc_middle::ty::Instance;
@@ -153,37 +153,38 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// Perform some simplifications to make the final coverage mappings
     /// slightly smaller.
     pub(crate) fn simplify_expressions(&mut self) {
-        // 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 `Operand::Zero`.
-        let mut zero_expressions = FxIndexSet::default();
-
-        // If an operand refers to an expression that is always zero, then
-        // that operand can be replaced with `Operand::Zero`.
-        let maybe_set_operand_to_zero =
-            |zero_expressions: &FxIndexSet<ExpressionId>, operand: &mut Operand| match &*operand {
-                Operand::Expression(id) if zero_expressions.contains(id) => {
-                    *operand = Operand::Zero;
+        // The set of expressions that were simplified to either `Zero` or a
+        // `Counter`. Other expressions that refer to these as operands
+        // can then also be simplified.
+        let mut simplified_expressions = FxHashMap::default();
+        type SimplifiedExpressions = FxHashMap<ExpressionId, Operand>;
+
+        // If an operand refers to an expression that has been simplified, then
+        // replace that operand with the simplified version.
+        let maybe_simplify_operand = |simplified_expressions: &SimplifiedExpressions,
+                                      operand: &mut Operand| {
+            if let Operand::Expression(id) = &*operand {
+                if let Some(simplified) = simplified_expressions.get(id) {
+                    *operand = *simplified;
                 }
-                _ => (),
-            };
+            }
+        };
 
         // For each expression, perform simplifications based on lower-numbered
-        // expressions, and then update the set of always-zero expressions if
+        // expressions, and then update the map of simplified expressions if
         // necessary.
         // (By construction, expressions can only refer to other expressions
         // that have lower IDs, so one simplification pass is sufficient.)
         for (id, maybe_expression) in self.expressions.iter_enumerated_mut() {
             let Some(expression) = maybe_expression else {
                 // If an expression is missing, it must have been optimized away,
                 // so any operand that refers to it can be replaced with zero.
-                zero_expressions.insert(id);
+                simplified_expressions.insert(id, Operand::Zero);
                 continue;
             };
 
-            maybe_set_operand_to_zero(&zero_expressions, &mut expression.lhs);
-            maybe_set_operand_to_zero(&zero_expressions, &mut expression.rhs);
+            maybe_simplify_operand(&simplified_expressions, &mut expression.lhs);
+            maybe_simplify_operand(&simplified_expressions, &mut expression.rhs);
 
             // Coverage counter values cannot be negative, so if an expression
             // involves subtraction from zero, assume that its RHS must also be zero.
@@ -192,18 +193,29 @@ impl<'tcx> FunctionCoverage<'tcx> {
                 expression.rhs = Operand::Zero;
             }
 
-            // After the above simplifications, if both operands are zero, then
-            // we know that this expression is always zero too.
-            if let Expression { lhs: Operand::Zero, rhs: Operand::Zero, .. } = expression {
-                zero_expressions.insert(id);
+            // After the above simplifications, if the right hand operand is zero,
+            // we can replace the expression by its left hand side.
+            if let Expression { lhs, rhs: Operand::Zero, .. } = expression {
+                simplified_expressions.insert(id, *lhs);
             }
         }
 
         // Do the same simplification for each side of a branch region.
         for branch in &mut self.branches {
-            maybe_set_operand_to_zero(&zero_expressions, &mut branch.true_);
-            maybe_set_operand_to_zero(&zero_expressions, &mut branch.false_);
+            maybe_simplify_operand(&simplified_expressions, &mut branch.true_);
+            maybe_simplify_operand(&simplified_expressions, &mut branch.false_);
+        }
+    }
+
+    /// This will further simplify any `Operand::Expression`,
+    /// "inlining" the left hand side operand if the right hand side is `Zero`.
+    fn simplified_expression(&self, id: ExpressionId) -> Counter {
+        if let Some(expr) = &self.expressions[id] {
+            if expr.rhs == Operand::Zero {
+                return Counter::from_operand(expr.lhs);
+            }
         }
+        Counter::expression(id)
     }
 
     /// Return the source hash, generated from the HIR node structure, and used to indicate whether
@@ -246,7 +258,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
             self.expressions.iter_enumerated().filter_map(|(id, expression)| {
                 let region = expression.as_ref()?.region.as_ref()?;
                 Some(CoverageCounterAndRegion {
-                    kind: CoverageCounterKind::Counter(Counter::expression(id)),
+                    kind: CoverageCounterKind::Counter(self.simplified_expression(id)),
                     region,
                 })
             });

tests/coverage-map/status-quo/async2.cov-map

@@ -40,15 +40,13 @@ Number of file 0 mappings: 1
 - Code(Counter(0)) at (prev + 31, 36) to (start + 2, 2)
 
 Function name: async2::executor::block_on::<async2::async_func::{closure#0}>
-Raw bytes (51): 0x[01, 01, 05, 0b, 05, 01, 05, 01, 05, 02, 00, 02, 00, 07, 01, 33, 05, 0a, 36, 20, 05, 02, 0d, 14, 00, 24, 02, 00, 20, 00, 23, 0b, 00, 27, 00, 49, 0f, 01, 17, 00, 1a, 05, 01, 0e, 00, 0f, 13, 02, 05, 00, 06]
+Raw bytes (47): 0x[01, 01, 03, 0b, 05, 01, 05, 01, 05, 07, 01, 33, 05, 0a, 36, 20, 05, 02, 0d, 14, 00, 24, 02, 00, 20, 00, 23, 0b, 00, 27, 00, 49, 02, 01, 17, 00, 1a, 05, 01, 0e, 00, 0f, 02, 02, 05, 00, 06]
 Number of files: 1
 - file 0 => global file 1
-Number of expressions: 5
+Number of expressions: 3
 - expression 0 operands: lhs = Expression(2, Add), rhs = Counter(1)
 - expression 1 operands: lhs = Counter(0), rhs = Counter(1)
 - expression 2 operands: lhs = Counter(0), rhs = Counter(1)
-- expression 3 operands: lhs = Expression(0, Sub), rhs = Zero
-- expression 4 operands: lhs = Expression(0, Sub), rhs = Zero
 Number of file 0 mappings: 7
 - Code(Counter(0)) at (prev + 51, 5) to (start + 10, 54)
 - Branch { true: Counter(1), false: Expression(0, Sub) } at (prev + 13, 20) to (start + 0, 36)
@@ -57,22 +55,20 @@ Number of file 0 mappings: 7
     = ((c0 + c1) - c1)
 - Code(Expression(2, Add)) at (prev + 0, 39) to (start + 0, 73)
     = (c0 + c1)
-- Code(Expression(3, Add)) at (prev + 1, 23) to (start + 0, 26)
-    = (((c0 + c1) - c1) + Zero)
+- Code(Expression(0, Sub)) at (prev + 1, 23) to (start + 0, 26)
+    = ((c0 + c1) - c1)
 - Code(Counter(1)) at (prev + 1, 14) to (start + 0, 15)
-- Code(Expression(4, Add)) at (prev + 2, 5) to (start + 0, 6)
-    = (((c0 + c1) - c1) + Zero)
+- Code(Expression(0, Sub)) at (prev + 2, 5) to (start + 0, 6)
+    = ((c0 + c1) - c1)
 
 Function name: async2::executor::block_on::<async2::async_func_just_println::{closure#0}>
-Raw bytes (51): 0x[01, 01, 05, 0b, 05, 01, 05, 01, 05, 02, 00, 02, 00, 07, 01, 33, 05, 0a, 36, 20, 05, 02, 0d, 14, 00, 24, 02, 00, 20, 00, 23, 0b, 00, 27, 00, 49, 0f, 01, 17, 00, 1a, 05, 01, 0e, 00, 0f, 13, 02, 05, 00, 06]
+Raw bytes (47): 0x[01, 01, 03, 0b, 05, 01, 05, 01, 05, 07, 01, 33, 05, 0a, 36, 20, 05, 02, 0d, 14, 00, 24, 02, 00, 20, 00, 23, 0b, 00, 27, 00, 49, 02, 01, 17, 00, 1a, 05, 01, 0e, 00, 0f, 02, 02, 05, 00, 06]
 Number of files: 1
 - file 0 => global file 1
-Number of expressions: 5
+Number of expressions: 3
 - expression 0 operands: lhs = Expression(2, Add), rhs = Counter(1)
 - expression 1 operands: lhs = Counter(0), rhs = Counter(1)
 - expression 2 operands: lhs = Counter(0), rhs = Counter(1)
-- expression 3 operands: lhs = Expression(0, Sub), rhs = Zero
-- expression 4 operands: lhs = Expression(0, Sub), rhs = Zero
 Number of file 0 mappings: 7
 - Code(Counter(0)) at (prev + 51, 5) to (start + 10, 54)
 - Branch { true: Counter(1), false: Expression(0, Sub) } at (prev + 13, 20) to (start + 0, 36)
@@ -81,11 +77,11 @@ Number of file 0 mappings: 7
     = ((c0 + c1) - c1)
 - Code(Expression(2, Add)) at (prev + 0, 39) to (start + 0, 73)
     = (c0 + c1)
-- Code(Expression(3, Add)) at (prev + 1, 23) to (start + 0, 26)
-    = (((c0 + c1) - c1) + Zero)
+- Code(Expression(0, Sub)) at (prev + 1, 23) to (start + 0, 26)
+    = ((c0 + c1) - c1)
 - Code(Counter(1)) at (prev + 1, 14) to (start + 0, 15)
-- Code(Expression(4, Add)) at (prev + 2, 5) to (start + 0, 6)
-    = (((c0 + c1) - c1) + Zero)
+- Code(Expression(0, Sub)) at (prev + 2, 5) to (start + 0, 6)
+    = ((c0 + c1) - c1)
 
 Function name: async2::executor::block_on::VTABLE::{closure#0}
 Raw bytes (9): 0x[01, 01, 00, 01, 01, 37, 11, 00, 33]
@@ -128,16 +124,14 @@ Number of file 0 mappings: 1
 - Code(Counter(0)) at (prev + 35, 1) to (start + 7, 2)
 
 Function name: async2::non_async_func
-Raw bytes (33): 0x[01, 01, 01, 05, 00, 05, 01, 09, 01, 03, 09, 20, 05, 00, 03, 08, 00, 09, 05, 00, 0a, 02, 06, 00, 02, 06, 00, 07, 03, 01, 01, 00, 02]
+Raw bytes (31): 0x[01, 01, 00, 05, 01, 09, 01, 03, 09, 20, 05, 00, 03, 08, 00, 09, 05, 00, 0a, 02, 06, 00, 02, 06, 00, 07, 05, 01, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
-Number of expressions: 1
-- expression 0 operands: lhs = Counter(1), rhs = Zero
+Number of expressions: 0
 Number of file 0 mappings: 5
 - Code(Counter(0)) at (prev + 9, 1) to (start + 3, 9)
 - Branch { true: Counter(1), false: Zero } at (prev + 3, 8) to (start + 0, 9)
 - Code(Counter(1)) at (prev + 0, 10) to (start + 2, 6)
 - Code(Zero) at (prev + 2, 6) to (start + 0, 7)
-- Code(Expression(0, Add)) at (prev + 1, 1) to (start + 0, 2)
-    = (c1 + Zero)
+- Code(Counter(1)) at (prev + 1, 1) to (start + 0, 2)