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_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,6 +461,9 @@ 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)?;
     }
@@ -468,6 +471,25 @@ pub fn write_mir_intro<'tcx>(
     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,