coverage: Initial support for branch coverage instrumentation

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,
     };