coverage: Data structures for recording branch info during MIR building

6 files changed, 94 insertions, 4 deletions

diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
index ed3a73af1ce..2c51aec3a76 100644
--- a/compiler/rustc_middle/src/mir/coverage.rs
+++ b/compiler/rustc_middle/src/mir/coverage.rs
@@ -2,7 +2,7 @@
 
 use rustc_index::IndexVec;
 use rustc_macros::HashStable;
-use rustc_span::Symbol;
+use rustc_span::{Span, Symbol};
 
 use std::fmt::{self, Debug, Formatter};
 
@@ -93,7 +93,7 @@ pub enum CoverageKind {
     SpanMarker,
 
     /// Marks its enclosing basic block with an ID that can be referred to by
-    /// other data in the MIR body.
+    /// side data in [`BranchInfo`].
     ///
     /// Has no effect during codegen.
     BlockMarker { id: BlockMarkerId },
@@ -218,3 +218,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,
+}
diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index b71c614dc4f..d57ffc0f8b5 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/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();
diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
index 8ae65f3832f..94751c44761 100644
--- a/compiler/rustc_middle/src/mir/pretty.rs
+++ b/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,
diff --git a/compiler/rustc_mir_build/src/build/coverageinfo.rs b/compiler/rustc_mir_build/src/build/coverageinfo.rs
new file mode 100644
index 00000000000..3fdf48f54c9
--- /dev/null
+++ b/compiler/rustc_mir_build/src/build/coverageinfo.rs
@@ -0,0 +1,32 @@
+use rustc_middle::mir;
+use rustc_middle::mir::coverage::BranchSpan;
+use rustc_middle::ty::TyCtxt;
+use rustc_span::def_id::LocalDefId;
+
+pub(crate) struct BranchInfoBuilder {
+    num_block_markers: usize,
+    branch_spans: Vec<BranchSpan>,
+}
+
+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 { num_block_markers: 0, branch_spans: vec![] })
+        } else {
+            None
+        }
+    }
+
+    pub(crate) fn into_done(self) -> Option<Box<mir::coverage::BranchInfo>> {
+        let Self { 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 }))
+    }
+}
diff --git a/compiler/rustc_mir_build/src/build/custom/mod.rs b/compiler/rustc_mir_build/src/build/custom/mod.rs
index c2bff9084c6..288b787798b 100644
--- a/compiler/rustc_mir_build/src/build/custom/mod.rs
+++ b/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,
     };
 
diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
index 45954bdb114..2fe5d3fb5e0 100644
--- a/compiler/rustc_mir_build/src/build/mod.rs
+++ b/compiler/rustc_mir_build/src/build/mod.rs
@@ -234,6 +234,10 @@ struct Builder<'a, 'tcx> {
     // the root (most of them do) and saves us from retracing many sub-paths
     // many times, and rechecking many nodes.
     lint_level_roots_cache: GrowableBitSet<hir::ItemLocalId>,
+
+    /// Collects additional coverage information during MIR building.
+    /// Only present if branch coverage is enabled and this function is eligible.
+    coverage_branch_info: Option<coverageinfo::BranchInfoBuilder>,
 }
 
 type CaptureMap<'tcx> = SortedIndexMultiMap<usize, hir::HirId, Capture<'tcx>>;
@@ -807,6 +811,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             unit_temp: None,
             var_debug_info: vec![],
             lint_level_roots_cache: GrowableBitSet::new_empty(),
+            coverage_branch_info: coverageinfo::BranchInfoBuilder::new_if_enabled(tcx, def),
         };
 
         assert_eq!(builder.cfg.start_new_block(), START_BLOCK);
@@ -826,7 +831,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
         }
 
-        Body::new(
+        let mut body = Body::new(
             MirSource::item(self.def_id.to_def_id()),
             self.cfg.basic_blocks,
             self.source_scopes,
@@ -837,7 +842,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             self.fn_span,
             self.coroutine,
             None,
-        )
+        );
+        body.coverage_branch_info = self.coverage_branch_info.and_then(|b| b.into_done());
+        body
     }
 
     fn insert_upvar_arg(&mut self) {
@@ -1111,6 +1118,7 @@ pub(crate) fn parse_float_into_scalar(
 
 mod block;
 mod cfg;
+mod coverageinfo;
 mod custom;
 mod expr;
 mod matches;