Rollup merge of #122322 - Zalathar:branch, r=oli-obk

coverage: Initial support for branch coverage instrumentation

(This is a review-ready version of the changes that were drafted in #118305.)

This PR adds support for branch coverage instrumentation, gated behind the unstable flag value `-Zcoverage-options=branch`. (Coverage instrumentation must also be enabled with `-Cinstrument-coverage`.)

During THIR-to-MIR lowering (MIR building), if branch coverage is enabled, we collect additional information about branch conditions and their corresponding then/else blocks. We inject special marker statements into those blocks, so that the `InstrumentCoverage` MIR pass can reliably identify them even after the initially-built MIR has been simplified and renumbered.

The rest of the changes are mostly just plumbing needed to gather up the information that was collected during MIR building, and include it in the coverage metadata that we embed in the final binary.

Note that `llvm-cov show` doesn't print branch coverage information in its source views by default; that needs to be explicitly enabled with `--show-branches=count` or similar.

---

The current implementation doesn't have any support for instrumenting `if let` or let-chains. I think it's still useful without that, and adding it would be non-trivial, so I'm happy to leave that for future work.

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -14,7 +14,6 @@ use self::spans::{BcbMapping, BcbMappingKind, CoverageSpans};
 use crate::MirPass;
 
 use rustc_middle::hir;
-use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::coverage::*;
 use rustc_middle::mir::{
     self, BasicBlock, BasicBlockData, Coverage, SourceInfo, Statement, StatementKind, Terminator,
@@ -44,7 +43,7 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
 
         let def_id = mir_source.def_id().expect_local();
 
-        if !is_eligible_for_coverage(tcx, def_id) {
+        if !tcx.is_eligible_for_coverage(def_id) {
             trace!("InstrumentCoverage skipped for {def_id:?} (not eligible)");
             return;
         }
@@ -140,6 +139,10 @@ fn create_mappings<'tcx>(
         .filter_map(|&BcbMapping { kind: bcb_mapping_kind, span }| {
             let kind = match bcb_mapping_kind {
                 BcbMappingKind::Code(bcb) => MappingKind::Code(term_for_bcb(bcb)),
+                BcbMappingKind::Branch { true_bcb, false_bcb } => MappingKind::Branch {
+                    true_term: term_for_bcb(true_bcb),
+                    false_term: term_for_bcb(false_bcb),
+                },
             };
             let code_region = make_code_region(source_map, file_name, span, body_span)?;
             Some(Mapping { kind, code_region })
@@ -349,37 +352,6 @@ fn check_code_region(code_region: CodeRegion) -> Option<CodeRegion> {
     }
 }
 
-fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
-    // Only instrument functions, methods, and closures (not constants since they are evaluated
-    // at compile time by Miri).
-    // FIXME(#73156): Handle source code coverage in const eval, but note, if and when const
-    // expressions get coverage spans, we will probably have to "carve out" space for const
-    // expressions from coverage spans in enclosing MIR's, like we do for closures. (That might
-    // be tricky if const expressions have no corresponding statements in the enclosing MIR.
-    // Closures are carved out by their initial `Assign` statement.)
-    if !tcx.def_kind(def_id).is_fn_like() {
-        trace!("InstrumentCoverage skipped for {def_id:?} (not an fn-like)");
-        return false;
-    }
-
-    // Don't instrument functions with `#[automatically_derived]` on their
-    // enclosing impl block, on the assumption that most users won't care about
-    // coverage for derived impls.
-    if let Some(impl_of) = tcx.impl_of_method(def_id.to_def_id())
-        && tcx.is_automatically_derived(impl_of)
-    {
-        trace!("InstrumentCoverage skipped for {def_id:?} (automatically derived)");
-        return false;
-    }
-
-    if tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
-        trace!("InstrumentCoverage skipped for {def_id:?} (`#[coverage(off)]`)");
-        return false;
-    }
-
-    true
-}
-
 /// Function information extracted from HIR by the coverage instrumentor.
 #[derive(Debug)]
 struct ExtractedHirInfo {

compiler/rustc_mir_transform/src/coverage/query.rs

@@ -1,14 +1,49 @@
-use super::*;
-
 use rustc_data_structures::captures::Captures;
-use rustc_middle::mir::coverage::*;
-use rustc_middle::mir::{Body, CoverageIdsInfo};
-use rustc_middle::query::Providers;
-use rustc_middle::ty::{self};
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use rustc_middle::mir::coverage::{CounterId, CoverageKind};
+use rustc_middle::mir::{Body, Coverage, CoverageIdsInfo, Statement, StatementKind};
+use rustc_middle::query::TyCtxtAt;
+use rustc_middle::ty::{self, TyCtxt};
+use rustc_middle::util::Providers;
+use rustc_span::def_id::LocalDefId;
 
-/// A `query` provider for retrieving coverage information injected into MIR.
+/// Registers query/hook implementations related to coverage.
 pub(crate) fn provide(providers: &mut Providers) {
-    providers.coverage_ids_info = |tcx, def_id| coverage_ids_info(tcx, def_id);
+    providers.hooks.is_eligible_for_coverage =
+        |TyCtxtAt { tcx, .. }, def_id| is_eligible_for_coverage(tcx, def_id);
+    providers.queries.coverage_ids_info = coverage_ids_info;
+}
+
+/// Hook implementation for [`TyCtxt::is_eligible_for_coverage`].
+fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
+    // Only instrument functions, methods, and closures (not constants since they are evaluated
+    // at compile time by Miri).
+    // FIXME(#73156): Handle source code coverage in const eval, but note, if and when const
+    // expressions get coverage spans, we will probably have to "carve out" space for const
+    // expressions from coverage spans in enclosing MIR's, like we do for closures. (That might
+    // be tricky if const expressions have no corresponding statements in the enclosing MIR.
+    // Closures are carved out by their initial `Assign` statement.)
+    if !tcx.def_kind(def_id).is_fn_like() {
+        trace!("InstrumentCoverage skipped for {def_id:?} (not an fn-like)");
+        return false;
+    }
+
+    // Don't instrument functions with `#[automatically_derived]` on their
+    // enclosing impl block, on the assumption that most users won't care about
+    // coverage for derived impls.
+    if let Some(impl_of) = tcx.impl_of_method(def_id.to_def_id())
+        && tcx.is_automatically_derived(impl_of)
+    {
+        trace!("InstrumentCoverage skipped for {def_id:?} (automatically derived)");
+        return false;
+    }
+
+    if tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
+        trace!("InstrumentCoverage skipped for {def_id:?} (`#[coverage(off)]`)");
+        return false;
+    }
+
+    true
 }
 
 /// Query implementation for `coverage_ids_info`.

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -13,6 +13,8 @@ mod from_mir;
 pub(super) enum BcbMappingKind {
     /// Associates an ordinary executable code span with its corresponding BCB.
     Code(BasicCoverageBlock),
+    /// Associates a branch span with BCBs for its true and false arms.
+    Branch { true_bcb: BasicCoverageBlock, false_bcb: BasicCoverageBlock },
 }
 
 #[derive(Debug)]
@@ -66,6 +68,12 @@ pub(super) fn generate_coverage_spans(
             // Each span produced by the generator represents an ordinary code region.
             BcbMapping { kind: BcbMappingKind::Code(bcb), span }
         }));
+
+        mappings.extend(from_mir::extract_branch_mappings(
+            mir_body,
+            hir_info.body_span,
+            basic_coverage_blocks,
+        ));
     }
 
     if mappings.is_empty() {
@@ -80,6 +88,10 @@ pub(super) fn generate_coverage_spans(
     for &BcbMapping { kind, span: _ } in &mappings {
         match kind {
             BcbMappingKind::Code(bcb) => insert(bcb),
+            BcbMappingKind::Branch { true_bcb, false_bcb } => {
+                insert(true_bcb);
+                insert(false_bcb);
+            }
         }
     }
 

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

@@ -1,7 +1,9 @@
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashSet;
+use rustc_index::IndexVec;
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
 use rustc_middle::mir::{
-    self, AggregateKind, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
+    self, AggregateKind, BasicBlock, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
     TerminatorKind,
 };
 use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
@@ -9,6 +11,7 @@ use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
 use crate::coverage::graph::{
     BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
 };
+use crate::coverage::spans::{BcbMapping, BcbMappingKind};
 use crate::coverage::ExtractedHirInfo;
 
 /// Traverses the MIR body to produce an initial collection of coverage-relevant
@@ -179,8 +182,6 @@ fn is_closure_like(statement: &Statement<'_>) -> bool {
 /// If the MIR `Statement` has a span contributive to computing coverage spans,
 /// return it; otherwise return `None`.
 fn filtered_statement_span(statement: &Statement<'_>) -> Option<Span> {
-    use mir::coverage::CoverageKind;
-
     match statement.kind {
         // These statements have spans that are often outside the scope of the executed source code
         // for their parent `BasicBlock`.
@@ -225,6 +226,11 @@ fn filtered_statement_span(statement: &Statement<'_>) -> Option<Span> {
             Some(statement.source_info.span)
         }
 
+        StatementKind::Coverage(box mir::Coverage {
+            // Block markers are used for branch coverage, so ignore them here.
+            kind: CoverageKind::BlockMarker {..}
+        }) => None,
+
         StatementKind::Coverage(box mir::Coverage {
             // These coverage statements should not exist prior to coverage instrumentation.
             kind: CoverageKind::CounterIncrement { .. } | CoverageKind::ExpressionUsed { .. }
@@ -358,3 +364,51 @@ impl SpanFromMir {
         Self { span, visible_macro, bcb, is_hole }
     }
 }
+
+pub(super) fn extract_branch_mappings(
+    mir_body: &mir::Body<'_>,
+    body_span: Span,
+    basic_coverage_blocks: &CoverageGraph,
+) -> Vec<BcbMapping> {
+    let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else {
+        return vec![];
+    };
+
+    let mut block_markers = IndexVec::<BlockMarkerId, Option<BasicBlock>>::from_elem_n(
+        None,
+        branch_info.num_block_markers,
+    );
+
+    // Fill out the mapping from block marker IDs to their enclosing blocks.
+    for (bb, data) in mir_body.basic_blocks.iter_enumerated() {
+        for statement in &data.statements {
+            if let StatementKind::Coverage(coverage) = &statement.kind
+                && let CoverageKind::BlockMarker { id } = coverage.kind
+            {
+                block_markers[id] = Some(bb);
+            }
+        }
+    }
+
+    branch_info
+        .branch_spans
+        .iter()
+        .filter_map(|&BranchSpan { span: raw_span, true_marker, false_marker }| {
+            // For now, ignore any branch span that was introduced by
+            // expansion. This makes things like assert macros less noisy.
+            if !raw_span.ctxt().outer_expn_data().is_root() {
+                return None;
+            }
+            let (span, _) = unexpand_into_body_span_with_visible_macro(raw_span, body_span)?;
+
+            let bcb_from_marker = |marker: BlockMarkerId| {
+                Some(basic_coverage_blocks.bcb_from_bb(block_markers[marker]?)?)
+            };
+
+            let true_bcb = bcb_from_marker(true_marker)?;
+            let false_bcb = bcb_from_marker(false_marker)?;
+
+            Some(BcbMapping { kind: BcbMappingKind::Branch { true_bcb, false_bcb }, span })
+        })
+        .collect::<Vec<_>>()
+}

compiler/rustc_mir_transform/src/lib.rs

@@ -37,8 +37,9 @@ use rustc_middle::mir::{
     LocalDecl, MirPass, MirPhase, Operand, Place, ProjectionElem, Promoted, RuntimePhase, Rvalue,
     SourceInfo, Statement, StatementKind, TerminatorKind, START_BLOCK,
 };
-use rustc_middle::query::Providers;
+use rustc_middle::query;
 use rustc_middle::ty::{self, TyCtxt, TypeVisitableExt};
+use rustc_middle::util::Providers;
 use rustc_span::{source_map::Spanned, sym, DUMMY_SP};
 use rustc_trait_selection::traits;
 
@@ -124,7 +125,7 @@ pub fn provide(providers: &mut Providers) {
     ffi_unwind_calls::provide(providers);
     shim::provide(providers);
     cross_crate_inline::provide(providers);
-    *providers = Providers {
+    providers.queries = query::Providers {
         mir_keys,
         mir_const,
         mir_const_qualif,
@@ -139,7 +140,7 @@ pub fn provide(providers: &mut Providers) {
         mir_inliner_callees: inline::cycle::mir_inliner_callees,
         promoted_mir,
         deduced_param_attrs: deduce_param_attrs::deduced_param_attrs,
-        ..*providers
+        ..providers.queries
     };
 }