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.

Author: matthiaskrgr

compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs

@@ -164,6 +164,15 @@ impl CounterMappingRegion {
                 end_line,
                 end_col,
             ),
+            MappingKind::Branch { true_term, false_term } => Self::branch_region(
+                Counter::from_term(true_term),
+                Counter::from_term(false_term),
+                local_file_id,
+                start_line,
+                start_col,
+                end_line,
+                end_col,
+            ),
         }
     }
 
@@ -188,9 +197,6 @@ impl CounterMappingRegion {
         }
     }
 
-    // This function might be used in the future; the LLVM API is still evolving, as is coverage
-    // support.
-    #[allow(dead_code)]
     pub(crate) fn branch_region(
         counter: Counter,
         false_counter: Counter,

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -88,7 +88,7 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
         match coverage.kind {
             // Marker statements have no effect during codegen,
             // so return early and don't create `func_coverage`.
-            CoverageKind::SpanMarker => return,
+            CoverageKind::SpanMarker | CoverageKind::BlockMarker { .. } => return,
             // Match exhaustively to ensure that newly-added kinds are classified correctly.
             CoverageKind::CounterIncrement { .. } | CoverageKind::ExpressionUsed { .. } => {}
         }
@@ -108,7 +108,7 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
 
         let Coverage { kind } = coverage;
         match *kind {
-            CoverageKind::SpanMarker => unreachable!(
+            CoverageKind::SpanMarker | CoverageKind::BlockMarker { .. } => unreachable!(
                 "unexpected marker statement {kind:?} should have caused an early return"
             ),
             CoverageKind::CounterIncrement { id } => {

compiler/rustc_middle/src/hooks/mod.rs

@@ -6,6 +6,7 @@
 use crate::mir;
 use crate::query::TyCtxtAt;
 use crate::ty::{Ty, TyCtxt};
+use rustc_span::def_id::LocalDefId;
 use rustc_span::DUMMY_SP;
 
 macro_rules! declare_hooks {
@@ -70,4 +71,10 @@ declare_hooks! {
 
     /// Getting a &core::panic::Location referring to a span.
     hook const_caller_location(file: rustc_span::Symbol, line: u32, col: u32) -> mir::ConstValue<'tcx>;
+
+    /// Returns `true` if this def is a function-like thing that is eligible for
+    /// coverage instrumentation under `-Cinstrument-coverage`.
+    ///
+    /// (Eligible functions might nevertheless be skipped for other reasons.)
+    hook is_eligible_for_coverage(key: LocalDefId) -> bool;
 }

compiler/rustc_middle/src/mir/coverage.rs

@@ -2,10 +2,19 @@
 
 use rustc_index::IndexVec;
 use rustc_macros::HashStable;
-use rustc_span::Symbol;
+use rustc_span::{Span, Symbol};
 
 use std::fmt::{self, Debug, Formatter};
 
+rustc_index::newtype_index! {
+    /// Used by [`CoverageKind::BlockMarker`] to mark blocks during THIR-to-MIR
+    /// lowering, so that those blocks can be identified later.
+    #[derive(HashStable)]
+    #[encodable]
+    #[debug_format = "BlockMarkerId({})"]
+    pub struct BlockMarkerId {}
+}
+
 rustc_index::newtype_index! {
     /// ID of a coverage counter. Values ascend from 0.
     ///
@@ -83,6 +92,12 @@ pub enum CoverageKind {
     /// codegen.
     SpanMarker,
 
+    /// Marks its enclosing basic block with an ID that can be referred to by
+    /// side data in [`BranchInfo`].
+    ///
+    /// Has no effect during codegen.
+    BlockMarker { id: BlockMarkerId },
+
     /// Marks the point in MIR control flow represented by a coverage counter.
     ///
     /// This is eventually lowered to `llvm.instrprof.increment` in LLVM IR.
@@ -107,6 +122,7 @@ impl Debug for CoverageKind {
         use CoverageKind::*;
         match self {
             SpanMarker => write!(fmt, "SpanMarker"),
+            BlockMarker { id } => write!(fmt, "BlockMarker({:?})", id.index()),
             CounterIncrement { id } => write!(fmt, "CounterIncrement({:?})", id.index()),
             ExpressionUsed { id } => write!(fmt, "ExpressionUsed({:?})", id.index()),
         }
@@ -163,14 +179,18 @@ pub struct Expression {
 pub enum MappingKind {
     /// Associates a normal region of code with a counter/expression/zero.
     Code(CovTerm),
+    /// Associates a branch region with separate counters for true and false.
+    Branch { true_term: CovTerm, false_term: CovTerm },
 }
 
 impl MappingKind {
     /// Iterator over all coverage terms in this mapping kind.
     pub fn terms(&self) -> impl Iterator<Item = CovTerm> {
-        let one = |a| std::iter::once(a);
+        let one = |a| std::iter::once(a).chain(None);
+        let two = |a, b| std::iter::once(a).chain(Some(b));
         match *self {
             Self::Code(term) => one(term),
+            Self::Branch { true_term, false_term } => two(true_term, false_term),
         }
     }
 
@@ -179,6 +199,9 @@ impl MappingKind {
     pub fn map_terms(&self, map_fn: impl Fn(CovTerm) -> CovTerm) -> Self {
         match *self {
             Self::Code(term) => Self::Code(map_fn(term)),
+            Self::Branch { true_term, false_term } => {
+                Self::Branch { true_term: map_fn(true_term), false_term: map_fn(false_term) }
+            }
         }
     }
 }
@@ -202,3 +225,22 @@ pub struct FunctionCoverageInfo {
     pub expressions: IndexVec<ExpressionId, Expression>,
     pub mappings: Vec<Mapping>,
 }
+
+/// Branch information recorded during THIR-to-MIR lowering, and stored in MIR.
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+pub struct BranchInfo {
+    /// 1 more than the highest-numbered [`CoverageKind::BlockMarker`] that was
+    /// injected into the MIR body. This makes it possible to allocate per-ID
+    /// data structures without having to scan the entire body first.
+    pub num_block_markers: usize,
+    pub branch_spans: Vec<BranchSpan>,
+}
+
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+pub struct BranchSpan {
+    pub span: Span,
+    pub true_marker: BlockMarkerId,
+    pub false_marker: BlockMarkerId,
+}

compiler/rustc_middle/src/mir/mod.rs

@@ -403,6 +403,12 @@ pub struct Body<'tcx> {
 
     pub tainted_by_errors: Option<ErrorGuaranteed>,
 
+    /// Branch coverage information collected during MIR building, to be used by
+    /// the `InstrumentCoverage` pass.
+    ///
+    /// Only present if branch coverage is enabled and this function is eligible.
+    pub coverage_branch_info: Option<Box<coverage::BranchInfo>>,
+
     /// Per-function coverage information added by the `InstrumentCoverage`
     /// pass, to be used in conjunction with the coverage statements injected
     /// into this body's blocks.
@@ -450,6 +456,7 @@ impl<'tcx> Body<'tcx> {
             is_polymorphic: false,
             injection_phase: None,
             tainted_by_errors,
+            coverage_branch_info: None,
             function_coverage_info: None,
         };
         body.is_polymorphic = body.has_non_region_param();
@@ -479,6 +486,7 @@ impl<'tcx> Body<'tcx> {
             is_polymorphic: false,
             injection_phase: None,
             tainted_by_errors: None,
+            coverage_branch_info: None,
             function_coverage_info: None,
         };
         body.is_polymorphic = body.has_non_region_param();

compiler/rustc_middle/src/mir/pretty.rs

@@ -461,13 +461,35 @@ pub fn write_mir_intro<'tcx>(
     // Add an empty line before the first block is printed.
     writeln!(w)?;
 
+    if let Some(branch_info) = &body.coverage_branch_info {
+        write_coverage_branch_info(branch_info, w)?;
+    }
     if let Some(function_coverage_info) = &body.function_coverage_info {
         write_function_coverage_info(function_coverage_info, w)?;
     }
 
     Ok(())
 }
 
+fn write_coverage_branch_info(
+    branch_info: &coverage::BranchInfo,
+    w: &mut dyn io::Write,
+) -> io::Result<()> {
+    let coverage::BranchInfo { branch_spans, .. } = branch_info;
+
+    for coverage::BranchSpan { span, true_marker, false_marker } in branch_spans {
+        writeln!(
+            w,
+            "{INDENT}coverage branch {{ true: {true_marker:?}, false: {false_marker:?} }} => {span:?}",
+        )?;
+    }
+    if !branch_spans.is_empty() {
+        writeln!(w)?;
+    }
+
+    Ok(())
+}
+
 fn write_function_coverage_info(
     function_coverage_info: &coverage::FunctionCoverageInfo,
     w: &mut dyn io::Write,

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -405,6 +405,7 @@ TrivialTypeTraversalImpls! {
     ::rustc_hir::HirId,
     ::rustc_hir::MatchSource,
     ::rustc_target::asm::InlineAsmRegOrRegClass,
+    crate::mir::coverage::BlockMarkerId,
     crate::mir::coverage::CounterId,
     crate::mir::coverage::ExpressionId,
     crate::mir::Local,

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -0,0 +1,148 @@
+use std::assert_matches::assert_matches;
+use std::collections::hash_map::Entry;
+
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
+use rustc_middle::mir::{self, BasicBlock, UnOp};
+use rustc_middle::thir::{ExprId, ExprKind, Thir};
+use rustc_middle::ty::TyCtxt;
+use rustc_span::def_id::LocalDefId;
+
+use crate::build::Builder;
+
+pub(crate) struct BranchInfoBuilder {
+    /// Maps condition expressions to their enclosing `!`, for better instrumentation.
+    nots: FxHashMap<ExprId, NotInfo>,
+
+    num_block_markers: usize,
+    branch_spans: Vec<BranchSpan>,
+}
+
+#[derive(Clone, Copy)]
+struct NotInfo {
+    /// When visiting the associated expression as a branch condition, treat this
+    /// enclosing `!` as the branch condition instead.
+    enclosing_not: ExprId,
+    /// True if the associated expression is nested within an odd number of `!`
+    /// expressions relative to `enclosing_not` (inclusive of `enclosing_not`).
+    is_flipped: bool,
+}
+
+impl BranchInfoBuilder {
+    /// Creates a new branch info builder, but only if branch coverage instrumentation
+    /// is enabled and `def_id` represents a function that is eligible for coverage.
+    pub(crate) fn new_if_enabled(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<Self> {
+        if tcx.sess.instrument_coverage_branch() && tcx.is_eligible_for_coverage(def_id) {
+            Some(Self { nots: FxHashMap::default(), num_block_markers: 0, branch_spans: vec![] })
+        } else {
+            None
+        }
+    }
+
+    /// Unary `!` expressions inside an `if` condition are lowered by lowering
+    /// their argument instead, and then reversing the then/else arms of that `if`.
+    ///
+    /// That's awkward for branch coverage instrumentation, so to work around that
+    /// we pre-emptively visit any affected `!` expressions, and record extra
+    /// information that [`Builder::visit_coverage_branch_condition`] can use to
+    /// synthesize branch instrumentation for the enclosing `!`.
+    pub(crate) fn visit_unary_not(&mut self, thir: &Thir<'_>, unary_not: ExprId) {
+        assert_matches!(thir[unary_not].kind, ExprKind::Unary { op: UnOp::Not, .. });
+
+        self.visit_with_not_info(
+            thir,
+            unary_not,
+            // Set `is_flipped: false` for the `!` itself, so that its enclosed
+            // expression will have `is_flipped: true`.
+            NotInfo { enclosing_not: unary_not, is_flipped: false },
+        );
+    }
+
+    fn visit_with_not_info(&mut self, thir: &Thir<'_>, expr_id: ExprId, not_info: NotInfo) {
+        match self.nots.entry(expr_id) {
+            // This expression has already been marked by an enclosing `!`.
+            Entry::Occupied(_) => return,
+            Entry::Vacant(entry) => entry.insert(not_info),
+        };
+
+        match thir[expr_id].kind {
+            ExprKind::Unary { op: UnOp::Not, arg } => {
+                // Invert the `is_flipped` flag for the contents of this `!`.
+                let not_info = NotInfo { is_flipped: !not_info.is_flipped, ..not_info };
+                self.visit_with_not_info(thir, arg, not_info);
+            }
+            ExprKind::Scope { value, .. } => self.visit_with_not_info(thir, value, not_info),
+            ExprKind::Use { source } => self.visit_with_not_info(thir, source, not_info),
+            // All other expressions (including `&&` and `||`) don't need any
+            // special handling of their contents, so stop visiting.
+            _ => {}
+        }
+    }
+
+    fn next_block_marker_id(&mut self) -> BlockMarkerId {
+        let id = BlockMarkerId::from_usize(self.num_block_markers);
+        self.num_block_markers += 1;
+        id
+    }
+
+    pub(crate) fn into_done(self) -> Option<Box<mir::coverage::BranchInfo>> {
+        let Self { nots: _, num_block_markers, branch_spans } = self;
+
+        if num_block_markers == 0 {
+            assert!(branch_spans.is_empty());
+            return None;
+        }
+
+        Some(Box::new(mir::coverage::BranchInfo { num_block_markers, branch_spans }))
+    }
+}
+
+impl Builder<'_, '_> {
+    /// If branch coverage is enabled, inject marker statements into `then_block`
+    /// and `else_block`, and record their IDs in the table of branch spans.
+    pub(crate) fn visit_coverage_branch_condition(
+        &mut self,
+        mut expr_id: ExprId,
+        mut then_block: BasicBlock,
+        mut else_block: BasicBlock,
+    ) {
+        // Bail out if branch coverage is not enabled for this function.
+        let Some(branch_info) = self.coverage_branch_info.as_ref() else { return };
+
+        // If this condition expression is nested within one or more `!` expressions,
+        // replace it with the enclosing `!` collected by `visit_unary_not`.
+        if let Some(&NotInfo { enclosing_not, is_flipped }) = branch_info.nots.get(&expr_id) {
+            expr_id = enclosing_not;
+            if is_flipped {
+                std::mem::swap(&mut then_block, &mut else_block);
+            }
+        }
+        let source_info = self.source_info(self.thir[expr_id].span);
+
+        // Now that we have `source_info`, we can upgrade to a &mut reference.
+        let branch_info = self.coverage_branch_info.as_mut().expect("upgrading & to &mut");
+
+        let mut inject_branch_marker = |block: BasicBlock| {
+            let id = branch_info.next_block_marker_id();
+
+            let marker_statement = mir::Statement {
+                source_info,
+                kind: mir::StatementKind::Coverage(Box::new(mir::Coverage {
+                    kind: CoverageKind::BlockMarker { id },
+                })),
+            };
+            self.cfg.push(block, marker_statement);
+
+            id
+        };
+
+        let true_marker = inject_branch_marker(then_block);
+        let false_marker = inject_branch_marker(else_block);
+
+        branch_info.branch_spans.push(BranchSpan {
+            span: source_info.span,
+            true_marker,
+            false_marker,
+        });
+    }
+}

compiler/rustc_mir_build/src/build/custom/mod.rs

@@ -60,6 +60,7 @@ pub(super) fn build_custom_mir<'tcx>(
         tainted_by_errors: None,
         injection_phase: None,
         pass_count: 0,
+        coverage_branch_info: None,
         function_coverage_info: None,
     };