Allow MIR pass UnreachableProp when coverage is enabled

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -1,9 +1,8 @@
-pub mod query;
-
 mod counters;
 mod graph;
+pub mod query;
+pub(crate) mod repair;
 mod spans;
-
 #[cfg(test)]
 mod tests;
 

compiler/rustc_mir_transform/src/coverage/repair.rs

@@ -0,0 +1,58 @@
+use rustc_middle::mir::coverage::{CounterId, CoverageKind};
+use rustc_middle::mir::{
+    self, Coverage, MirPass, SourceInfo, Statement, StatementKind, START_BLOCK,
+};
+use rustc_middle::ty::TyCtxt;
+use rustc_span::DUMMY_SP;
+
+/// If a function has been [instrumented for coverage](super::InstrumentCoverage),
+/// but MIR optimizations subsequently remove all of its [`CoverageKind::CounterIncrement`]
+/// statements (e.g. because bb0 is unreachable), then we won't generate any
+/// `llvm.instrprof.increment` intrinsics. LLVM will think the function is not
+/// instrumented, and it will disappear from coverage mappings and coverage reports.
+///
+/// This pass detects when that has happened, and re-inserts a dummy counter-increment
+/// statement so that LLVM knows to treat the function as instrumented.
+pub struct RepairCoverage;
+
+impl<'tcx> MirPass<'tcx> for RepairCoverage {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.instrument_coverage()
+    }
+
+    fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
+        // If a function wasn't instrumented for coverage in the first place,
+        // then there's no need to repair anything.
+        if body.function_coverage_info.is_none() {
+            return;
+        }
+
+        // If the body still contains one or more counter-increment statements,
+        // there's no need to repair anything.
+        let has_counter_increment = body
+            .basic_blocks
+            .iter()
+            .flat_map(|bb_data| &bb_data.statements)
+            .filter_map(|statement| match statement.kind {
+                StatementKind::Coverage(box ref coverage) => Some(coverage),
+                _ => None,
+            })
+            .any(|coverage| matches!(coverage.kind, CoverageKind::CounterIncrement { .. }));
+        if has_counter_increment {
+            return;
+        }
+
+        debug!(
+            "all counter-increments were removed after instrumentation; restoring one counter in {:?}",
+            body.source.def_id(),
+        );
+
+        let statement = Statement {
+            source_info: SourceInfo::outermost(DUMMY_SP),
+            kind: StatementKind::Coverage(Box::new(Coverage {
+                kind: CoverageKind::CounterIncrement { id: CounterId::START },
+            })),
+        };
+        body[START_BLOCK].statements.insert(0, statement);
+    }
+}

compiler/rustc_mir_transform/src/lib.rs

@@ -585,6 +585,9 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &large_enums::EnumSizeOpt { discrepancy: 128 },
             // Some cleanup necessary at least for LLVM and potentially other codegen backends.
             &add_call_guards::CriticalCallEdges,
+            // If MIR optimizations removed all coverage-increment statements
+            // from an instrumented function, add another one to avoid problems.
+            &coverage::repair::RepairCoverage,
             // Cleanup for human readability, off by default.
             &prettify::ReorderBasicBlocks,
             &prettify::ReorderLocals,

compiler/rustc_mir_transform/src/unreachable_prop.rs

@@ -13,11 +13,7 @@ pub struct UnreachablePropagation;
 impl MirPass<'_> for UnreachablePropagation {
     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
         // Enable only under -Zmir-opt-level=2 as this can make programs less debuggable.
-
-        // FIXME(#116171) Coverage gets confused by MIR passes that can remove all
-        // coverage statements from an instrumented function. This pass can be
-        // re-enabled when coverage codegen is robust against that happening.
-        sess.mir_opt_level() >= 2 && !sess.instrument_coverage()
+        sess.mir_opt_level() >= 2
     }
 
     fn run_pass<'tcx>(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {

tests/coverage-map/unreachable.rs

@@ -1,6 +1,6 @@
 #![feature(core_intrinsics)]
 #![feature(coverage_attribute)]
-// compile-flags: --edition=2021
+// compile-flags: --edition=2021 -Copt-level=2
 
 // <https://github.com/rust-lang/rust/issues/116171>
 // If we instrument a function for coverage, but all of its counter-increment

tests/run-coverage/unreachable.coverage

@@ -1,6 +1,6 @@
    LL|       |#![feature(core_intrinsics)]
    LL|       |#![feature(coverage_attribute)]
-   LL|       |// compile-flags: --edition=2021
+   LL|       |// compile-flags: --edition=2021 -Copt-level=2
    LL|       |
    LL|       |// <https://github.com/rust-lang/rust/issues/116171>
    LL|       |// If we instrument a function for coverage, but all of its counter-increment

tests/run-coverage/unreachable.rs

@@ -1,6 +1,6 @@
 #![feature(core_intrinsics)]
 #![feature(coverage_attribute)]
-// compile-flags: --edition=2021
+// compile-flags: --edition=2021 -Copt-level=2
 
 // <https://github.com/rust-lang/rust/issues/116171>
 // If we instrument a function for coverage, but all of its counter-increment