coverage: Include recorded branch info in coverage instrumentation/codegen

Author: Zalathar

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,
+            ),
         }
     }
 

compiler/rustc_middle/src/mir/coverage.rs

@@ -179,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),
         }
     }
 
@@ -195,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) }
+            }
         }
     }
 }

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -139,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 })

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,14 +1,17 @@
 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};
 
 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`.
@@ -363,3 +364,50 @@ 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(hir_branch_info) = mir_body.coverage_hir_branch_info.as_deref() else {
+        return vec![];
+    };
+
+    let mut block_markers = IndexVec::<BlockMarkerId, Option<BasicBlock>>::from_elem_n(
+        None,
+        hir_branch_info.num_block_markers,
+    );
+
+    for (bb, data) in mir_body.basic_blocks.iter_enumerated() {
+        for statement in &data.statements {
+            if let StatementKind::Coverage(coverage) = &statement.kind {
+                if let CoverageKind::BlockMarker { id } = coverage.kind {
+                    block_markers[id] = Some(bb);
+                }
+            }
+        }
+    }
+
+    hir_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<_>>()
+}