coverage: Clean up marker statements that aren't needed later

Some of the marker statements used by coverage are added during MIR building
for use by the InstrumentCoverage pass (during analysis), and are not needed
afterwards.

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -85,14 +85,6 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
 
         let bx = self;
 
-        match coverage.kind {
-            // Marker statements have no effect during codegen,
-            // so return early and don't create `func_coverage`.
-            CoverageKind::SpanMarker | CoverageKind::BlockMarker { .. } => return,
-            // Match exhaustively to ensure that newly-added kinds are classified correctly.
-            CoverageKind::CounterIncrement { .. } | CoverageKind::ExpressionUsed { .. } => {}
-        }
-
         let Some(function_coverage_info) =
             bx.tcx.instance_mir(instance.def).function_coverage_info.as_deref()
         else {
@@ -109,7 +101,7 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
         let Coverage { kind } = coverage;
         match *kind {
             CoverageKind::SpanMarker | CoverageKind::BlockMarker { .. } => unreachable!(
-                "unexpected marker statement {kind:?} should have caused an early return"
+                "marker statement {kind:?} should have been removed by CleanupPostBorrowck"
             ),
             CoverageKind::CounterIncrement { id } => {
                 func_coverage.mark_counter_id_seen(id);

compiler/rustc_middle/src/mir/coverage.rs

@@ -88,14 +88,13 @@ pub enum CoverageKind {
     /// Marks a span that might otherwise not be represented in MIR, so that
     /// coverage instrumentation can associate it with its enclosing block/BCB.
     ///
-    /// Only used by the `InstrumentCoverage` pass, and has no effect during
-    /// codegen.
+    /// Should be erased before codegen (at some point after `InstrumentCoverage`).
     SpanMarker,
 
     /// Marks its enclosing basic block with an ID that can be referred to by
     /// side data in [`BranchInfo`].
     ///
-    /// Has no effect during codegen.
+    /// Should be erased before codegen (at some point after `InstrumentCoverage`).
     BlockMarker { id: BlockMarkerId },
 
     /// Marks the point in MIR control flow represented by a coverage counter.

compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs

@@ -5,15 +5,20 @@
 //!   - [`AscribeUserType`]
 //!   - [`FakeRead`]
 //!   - [`Assign`] statements with a [`Fake`] borrow
+//!   - [`Coverage`] statements of kind [`BlockMarker`] or [`SpanMarker`]
 //!
 //! [`AscribeUserType`]: rustc_middle::mir::StatementKind::AscribeUserType
 //! [`Assign`]: rustc_middle::mir::StatementKind::Assign
 //! [`FakeRead`]: rustc_middle::mir::StatementKind::FakeRead
 //! [`Nop`]: rustc_middle::mir::StatementKind::Nop
 //! [`Fake`]: rustc_middle::mir::BorrowKind::Fake
+//! [`Coverage`]: rustc_middle::mir::StatementKind::Coverage
+//! [`BlockMarker`]: rustc_middle::mir::coverage::CoverageKind::BlockMarker
+//! [`SpanMarker`]: rustc_middle::mir::coverage::CoverageKind::SpanMarker
 
 use crate::MirPass;
-use rustc_middle::mir::{Body, BorrowKind, Rvalue, StatementKind, TerminatorKind};
+use rustc_middle::mir::coverage::CoverageKind;
+use rustc_middle::mir::{Body, BorrowKind, Coverage, Rvalue, StatementKind, TerminatorKind};
 use rustc_middle::ty::TyCtxt;
 
 pub struct CleanupPostBorrowck;
@@ -25,6 +30,12 @@ impl<'tcx> MirPass<'tcx> for CleanupPostBorrowck {
                 match statement.kind {
                     StatementKind::AscribeUserType(..)
                     | StatementKind::Assign(box (_, Rvalue::Ref(_, BorrowKind::Fake, _)))
+                    | StatementKind::Coverage(box Coverage {
+                        // These kinds of coverage statements are markers inserted during
+                        // MIR building, and are not needed after InstrumentCoverage.
+                        kind: CoverageKind::BlockMarker { .. } | CoverageKind::SpanMarker { .. },
+                        ..
+                    })
                     | StatementKind::FakeRead(..) => statement.make_nop(),
                     _ => (),
                 }

tests/mir-opt/instrument_coverage_cleanup.main.CleanupPostBorrowck.diff

@@ -0,0 +1,56 @@
+- // MIR for `main` before CleanupPostBorrowck
++ // MIR for `main` after CleanupPostBorrowck
+  
+  fn main() -> () {
+      let mut _0: ();
+      let mut _1: bool;
+  
+      coverage branch { true: BlockMarkerId(0), false: BlockMarkerId(1) } => /the/src/instrument_coverage_cleanup.rs:17:8: 17:36 (#0)
+  
+      coverage ExpressionId(0) => Expression { lhs: Counter(0), op: Subtract, rhs: Counter(1) };
+      coverage ExpressionId(1) => Expression { lhs: Counter(1), op: Add, rhs: Expression(0) };
+      coverage Code(Counter(0)) => /the/src/instrument_coverage_cleanup.rs:16:1 - 17:36;
+      coverage Code(Expression(0)) => /the/src/instrument_coverage_cleanup.rs:17:37 - 17:39;
+      coverage Code(Counter(1)) => /the/src/instrument_coverage_cleanup.rs:17:39 - 17:40;
+      coverage Code(Expression(1)) => /the/src/instrument_coverage_cleanup.rs:18:1 - 18:2;
+      coverage Branch { true_term: Expression(0), false_term: Counter(1) } => /the/src/instrument_coverage_cleanup.rs:17:8 - 17:36;
+  
+      bb0: {
+          Coverage::CounterIncrement(0);
+-         Coverage::SpanMarker;
++         nop;
+          StorageLive(_1);
+          _1 = std::hint::black_box::<bool>(const true) -> [return: bb1, unwind: bb5];
+      }
+  
+      bb1: {
+          switchInt(move _1) -> [0: bb3, otherwise: bb2];
+      }
+  
+      bb2: {
+          Coverage::CounterIncrement(1);
+-         Coverage::BlockMarker(1);
++         nop;
+          _0 = const ();
+          goto -> bb4;
+      }
+  
+      bb3: {
+          Coverage::ExpressionUsed(0);
+-         Coverage::BlockMarker(0);
++         nop;
+          _0 = const ();
+          goto -> bb4;
+      }
+  
+      bb4: {
+          Coverage::ExpressionUsed(1);
+          StorageDead(_1);
+          return;
+      }
+  
+      bb5 (cleanup): {
+          resume;
+      }
+  }
+  

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

@@ -0,0 +1,53 @@
+- // MIR for `main` before InstrumentCoverage
++ // MIR for `main` after InstrumentCoverage
+  
+  fn main() -> () {
+      let mut _0: ();
+      let mut _1: bool;
+  
+      coverage branch { true: BlockMarkerId(0), false: BlockMarkerId(1) } => /the/src/instrument_coverage_cleanup.rs:17:8: 17:36 (#0)
+  
++     coverage ExpressionId(0) => Expression { lhs: Counter(0), op: Subtract, rhs: Counter(1) };
++     coverage ExpressionId(1) => Expression { lhs: Counter(1), op: Add, rhs: Expression(0) };
++     coverage Code(Counter(0)) => /the/src/instrument_coverage_cleanup.rs:16:1 - 17:36;
++     coverage Code(Expression(0)) => /the/src/instrument_coverage_cleanup.rs:17:37 - 17:39;
++     coverage Code(Counter(1)) => /the/src/instrument_coverage_cleanup.rs:17:39 - 17:40;
++     coverage Code(Expression(1)) => /the/src/instrument_coverage_cleanup.rs:18:1 - 18:2;
++     coverage Branch { true_term: Expression(0), false_term: Counter(1) } => /the/src/instrument_coverage_cleanup.rs:17:8 - 17:36;
++ 
+      bb0: {
++         Coverage::CounterIncrement(0);
+          Coverage::SpanMarker;
+          StorageLive(_1);
+          _1 = std::hint::black_box::<bool>(const true) -> [return: bb1, unwind: bb5];
+      }
+  
+      bb1: {
+          switchInt(move _1) -> [0: bb3, otherwise: bb2];
+      }
+  
+      bb2: {
++         Coverage::CounterIncrement(1);
+          Coverage::BlockMarker(1);
+          _0 = const ();
+          goto -> bb4;
+      }
+  
+      bb3: {
++         Coverage::ExpressionUsed(0);
+          Coverage::BlockMarker(0);
+          _0 = const ();
+          goto -> bb4;
+      }
+  
+      bb4: {
++         Coverage::ExpressionUsed(1);
+          StorageDead(_1);
+          return;
+      }
+  
+      bb5 (cleanup): {
+          resume;
+      }
+  }
+  

tests/mir-opt/instrument_coverage_cleanup.rs

@@ -0,0 +1,24 @@
+// Test that CleanupPostBorrowck cleans up the marker statements that are
+// inserted during MIR building (after InstrumentCoverage is done with them),
+// but leaves the statements that were added by InstrumentCoverage.
+//
+// Removed statement kinds: BlockMarker, SpanMarker
+// Retained statement kinds: CounterIncrement, ExpressionUsed
+
+//@ unit-test: InstrumentCoverage
+//@ needs-profiler-support
+//@ ignore-windows (see ./instrument_coverage.rs)
+//@ compile-flags: -Cinstrument-coverage -Zcoverage-options=branch
+//@ compile-flags: --remap-path-prefix={{src-base}}=/the/src
+
+// EMIT_MIR instrument_coverage_cleanup.main.InstrumentCoverage.diff
+// EMIT_MIR instrument_coverage_cleanup.main.CleanupPostBorrowck.diff
+fn main() {
+    if !core::hint::black_box(true) {}
+}
+
+// CHECK-NOT: Coverage::BlockMarker
+// CHECK-NOT: Coverage::SpanMarker
+// CHECK:     Coverage::CounterIncrement
+// CHECK-NOT: Coverage::BlockMarker
+// CHECK-NOT: Coverage::SpanMarker