coverage: Emit branch mappings

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -4,7 +4,8 @@ use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_index::bit_set::BitSet;
 use rustc_middle::mir::coverage::{
-    CodeRegion, CounterId, CovTerm, Expression, ExpressionId, FunctionCoverageInfo, Mapping, Op,
+    BranchMapping, CodeRegion, CounterId, CovTerm, Expression, ExpressionId, FunctionCoverageInfo,
+    Mapping, Op,
 };
 use rustc_middle::ty::Instance;
 use rustc_span::Symbol;
@@ -192,7 +193,18 @@ impl<'tcx> FunctionCoverage<'tcx> {
 
     /// Returns an iterator over all filenames used by this function's mappings.
     pub(crate) fn all_file_names(&self) -> impl Iterator<Item = Symbol> + Captures<'_> {
-        self.function_coverage_info.mappings.iter().map(|mapping| mapping.code_region.file_name)
+        let mapping_file_names = self
+            .function_coverage_info
+            .mappings
+            .iter()
+            .map(|mapping| mapping.code_region.file_name);
+        let branch_mapping_file_names = self
+            .function_coverage_info
+            .branch_mappings
+            .iter()
+            .map(|branch_mapping| branch_mapping.code_region.file_name);
+
+        mapping_file_names.chain(branch_mapping_file_names)
     }
 
     /// Convert this function's coverage expression data into a form that can be
@@ -229,6 +241,17 @@ impl<'tcx> FunctionCoverage<'tcx> {
         })
     }
 
+    pub(crate) fn branch_regions(&self) -> impl Iterator<Item = (&CodeRegion, Counter, Counter)> {
+        self.function_coverage_info.branch_mappings.iter().map(move |branch_mapping| {
+            let &BranchMapping { ref code_region, true_arm_term, false_arm_term } = branch_mapping;
+            (
+                code_region,
+                self.counter_for_term(true_arm_term),
+                self.counter_for_term(false_arm_term),
+            )
+        })
+    }
+
     fn counter_for_term(&self, term: CovTerm) -> Counter {
         if is_zero_term(&self.counters_seen, &self.zero_expressions, term) {
             Counter::ZERO

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -263,6 +263,23 @@ fn encode_mappings_for_function(
         }
     }
 
+    for (region, true_counter, false_counter) in function_coverage.branch_regions() {
+        let CodeRegion { file_name, start_line, start_col, end_line, end_col } = *region;
+
+        let global_file_id = global_file_table.global_file_id_for_file_name(file_name);
+        let local_file_id = virtual_file_mapping.local_id_for_global(global_file_id);
+
+        mapping_regions.push(CounterMappingRegion::branch_region(
+            true_counter,
+            false_counter,
+            local_file_id.as_u32(),
+            start_line,
+            start_col,
+            end_line,
+            end_col,
+        ));
+    }
+
     // Encode the function's coverage mappings into a buffer.
     llvm::build_byte_buffer(|buffer| {
         coverageinfo::write_mapping_to_buffer(

compiler/rustc_middle/src/mir/coverage.rs

@@ -173,6 +173,15 @@ pub struct Mapping {
     pub term: CovTerm,
 }
 
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+pub struct BranchMapping {
+    pub code_region: CodeRegion,
+
+    pub true_arm_term: CovTerm,
+    pub false_arm_term: CovTerm,
+}
+
 /// Stores per-function coverage information attached to a `mir::Body`,
 /// to be used in conjunction with the individual coverage statements injected
 /// into the function's basic blocks.
@@ -184,6 +193,7 @@ pub struct FunctionCoverageInfo {
 
     pub expressions: IndexVec<ExpressionId, Expression>,
     pub mappings: Vec<Mapping>,
+    pub branch_mappings: Vec<BranchMapping>,
 }
 
 /// Coverage information captured from HIR/THIR during MIR building.

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -9,7 +9,7 @@ mod tests;
 
 use self::counters::{BcbCounter, CoverageCounters};
 use self::graph::CoverageGraph;
-use self::spans::CoverageSpans;
+use self::spans::{BcbBranchSpan, CoverageSpans};
 
 use crate::MirPass;
 
@@ -124,13 +124,15 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
         self.coverage_counters
             .make_bcb_counters(&self.basic_coverage_blocks, bcb_has_coverage_spans);
 
-        let mappings = self.create_mappings_and_inject_coverage_statements(&coverage_spans);
+        let (mappings, branch_mappings) =
+            self.create_mappings_and_inject_coverage_statements(&coverage_spans);
 
         self.mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo {
             function_source_hash: self.function_source_hash,
             num_counters: self.coverage_counters.num_counters(),
             expressions: self.coverage_counters.take_expressions(),
             mappings,
+            branch_mappings,
         }));
     }
 
@@ -140,7 +142,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
     fn create_mappings_and_inject_coverage_statements(
         &mut self,
         coverage_spans: &CoverageSpans,
-    ) -> Vec<Mapping> {
+    ) -> (Vec<Mapping>, Vec<BranchMapping>) {
         let source_map = self.tcx.sess.source_map();
         let body_span = self.body_span;
 
@@ -210,7 +212,17 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             inject_statement(self.mir_body, self.make_mir_coverage_kind(counter_kind), new_bb);
         }
 
-        mappings
+        let branch_mappings = coverage_spans
+            .branch_spans()
+            .iter()
+            .map(|&BcbBranchSpan { span, true_bcb, false_bcb }| BranchMapping {
+                code_region: make_code_region(source_map, file_name, span, body_span),
+                true_arm_term: self.coverage_counters.bcb_counter(true_bcb).unwrap().as_term(),
+                false_arm_term: self.coverage_counters.bcb_counter(false_bcb).unwrap().as_term(),
+            })
+            .collect::<Vec<_>>();
+
+        (mappings, branch_mappings)
     }
 
     fn make_mir_coverage_kind(&self, counter_kind: &BcbCounter) -> CoverageKind {

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -1,17 +1,28 @@
 use std::cell::OnceCell;
 
 use rustc_data_structures::graph::WithNumNodes;
+use rustc_index::bit_set::BitSet;
 use rustc_index::IndexVec;
-use rustc_middle::mir;
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
+use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol, DUMMY_SP};
 
 use super::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
 
 mod from_mir;
 
+pub(super) struct BcbBranchSpan {
+    pub(super) span: Span,
+    pub(super) true_bcb: BasicCoverageBlock,
+    pub(super) false_bcb: BasicCoverageBlock,
+}
+
 pub(super) struct CoverageSpans {
     /// Map from BCBs to their list of coverage spans.
     bcb_to_spans: IndexVec<BasicCoverageBlock, Vec<Span>>,
+
+    bcbs_with_branch_spans: BitSet<BasicCoverageBlock>,
+    branch_spans: Vec<BcbBranchSpan>,
 }
 
 impl CoverageSpans {
@@ -34,16 +45,57 @@ impl CoverageSpans {
             bcb_to_spans[bcb].push(span);
         }
 
-        Self { bcb_to_spans }
+        let mut bcbs_with_branch_spans = BitSet::new_empty(basic_coverage_blocks.num_nodes());
+        let mut branch_spans = Vec::new();
+
+        if let Some(coverage_hir_info) = mir_body.coverage_hir_info.as_deref() {
+            let mut block_markers = IndexVec::<BlockMarkerId, Option<BasicBlock>>::from_elem_n(
+                None,
+                coverage_hir_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);
+                        }
+                    }
+                }
+            }
+
+            // Assume that most branch spans will resolve successfully.
+            branch_spans.reserve(coverage_hir_info.branch_spans.len());
+            branch_spans.extend(coverage_hir_info.branch_spans.iter().filter_map(
+                |&BranchSpan { span, true_marker, false_marker }| {
+                    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)?;
+
+                    bcbs_with_branch_spans.insert(true_bcb);
+                    bcbs_with_branch_spans.insert(false_bcb);
+                    Some(BcbBranchSpan { span, true_bcb, false_bcb })
+                },
+            ));
+        }
+
+        Self { bcb_to_spans, bcbs_with_branch_spans, branch_spans }
     }
 
     pub(super) fn bcb_has_coverage_spans(&self, bcb: BasicCoverageBlock) -> bool {
-        !self.bcb_to_spans[bcb].is_empty()
+        !self.bcb_to_spans[bcb].is_empty() || self.bcbs_with_branch_spans.contains(bcb)
     }
 
     pub(super) fn spans_for_bcb(&self, bcb: BasicCoverageBlock) -> &[Span] {
         &self.bcb_to_spans[bcb]
     }
+
+    pub(super) fn branch_spans(&self) -> &[BcbBranchSpan] {
+        &self.branch_spans
+    }
 }
 
 /// A BCB is deconstructed into one or more `Span`s. Each `Span` maps to a `CoverageSpan` that