coverage: Data structures for recording branch info during MIR building

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -100,6 +100,8 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
 
         let Coverage { kind } = coverage;
         match *kind {
+            // Block markers have no effect during codegen.
+            CoverageKind::BlockMarker { .. } => {}
             CoverageKind::CounterIncrement { id } => {
                 func_coverage.mark_counter_id_seen(id);
                 // We need to explicitly drop the `RefMut` before calling into `instrprof_increment`,

compiler/rustc_middle/src/mir/coverage.rs

@@ -76,13 +76,19 @@ impl Debug for CovTerm {
 
 #[derive(Clone, PartialEq, TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub enum CoverageKind {
+    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.
     ///
     /// If this statement does not survive MIR optimizations, any mappings that
     /// refer to this counter can have those references simplified to zero.
-    CounterIncrement { id: CounterId },
+    CounterIncrement {
+        id: CounterId,
+    },
 
     /// Marks the point in MIR control-flow represented by a coverage expression.
     ///
@@ -92,13 +98,16 @@ pub enum CoverageKind {
     /// (This is only inserted for expression IDs that are directly used by
     /// mappings. Intermediate expressions with no direct mappings are
     /// retained/zeroed based on whether they are transitively used.)
-    ExpressionUsed { id: ExpressionId },
+    ExpressionUsed {
+        id: ExpressionId,
+    },
 }
 
 impl Debug for CoverageKind {
     fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
         use CoverageKind::*;
         match self {
+            BlockMarker { id } => write!(fmt, "BlockMarker({:?})", id.index()),
             CounterIncrement { id } => write!(fmt, "CounterIncrement({:?})", id.index()),
             ExpressionUsed { id } => write!(fmt, "ExpressionUsed({:?})", id.index()),
         }
@@ -187,4 +196,22 @@ pub struct HirInfo {
     pub function_source_hash: u64,
     pub fn_sig_span: Span,
     pub body_span: Span,
+
+    pub num_block_markers: usize,
+    pub branch_spans: Vec<BranchSpan>,
+}
+
+rustc_index::newtype_index! {
+    #[derive(HashStable)]
+    #[encodable]
+    #[debug_format = "BlockMarkerId({})"]
+    pub struct BlockMarkerId {}
+}
+
+#[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/pretty.rs

@@ -493,13 +493,33 @@ pub fn write_mir_intro<'tcx>(
     // Add an empty line before the first block is printed.
     writeln!(w)?;
 
+    if let Some(coverage_hir_info) = &body.coverage_hir_info {
+        write_coverage_hir_info(coverage_hir_info, w)?;
+    }
     if let Some(function_coverage_info) = &body.function_coverage_info {
         write_function_coverage_info(function_coverage_info, w)?;
     }
 
     Ok(())
 }
 
+fn write_coverage_hir_info(
+    coverage_hir_info: &coverage::HirInfo,
+    w: &mut dyn io::Write,
+) -> io::Result<()> {
+    let coverage::HirInfo { branch_spans, .. } = coverage_hir_info;
+
+    for coverage::BranchSpan { span, true_marker, false_marker } in branch_spans {
+        writeln!(
+            w,
+            "{INDENT}coverage branch {{ true: {true_marker:?}, false: {false_marker:?} }} => {span:?}",
+        )?;
+    }
+    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

@@ -397,6 +397,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

@@ -1,26 +1,32 @@
 use rustc_middle::hir;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir;
+use rustc_middle::mir::coverage::BranchSpan;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::LocalDefId;
 use rustc_span::{ExpnKind, Span};
 
 pub(crate) struct HirInfoBuilder {
-    // Currently there is no extra state tracked here.
+    pub(crate) branch_coverage_enabled: bool,
+    pub(crate) num_block_markers: usize,
+    pub(crate) branch_spans: Vec<BranchSpan>,
 }
 
 impl HirInfoBuilder {
     pub(crate) fn new_if_enabled_and_eligible(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<Self> {
         if tcx.sess.instrument_coverage() && is_eligible_for_coverage(tcx, def_id) {
-            Some(Self {})
+            Some(Self {
+                branch_coverage_enabled: tcx.sess.instrument_coverage_branch(),
+                num_block_markers: 0,
+                branch_spans: Vec::new(),
+            })
         } else {
             None
         }
     }
 
     pub(crate) fn finish(self, tcx: TyCtxt<'_>, def_id: LocalDefId) -> Box<mir::coverage::HirInfo> {
-        let Self {} = self;
-        Box::new(make_coverage_hir_info(tcx, def_id))
+        Box::new(make_coverage_hir_info(tcx, def_id, self))
     }
 }
 
@@ -46,7 +52,11 @@ fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
     true
 }
 
-fn make_coverage_hir_info(tcx: TyCtxt<'_>, def_id: LocalDefId) -> mir::coverage::HirInfo {
+fn make_coverage_hir_info(
+    tcx: TyCtxt<'_>,
+    def_id: LocalDefId,
+    builder: HirInfoBuilder,
+) -> mir::coverage::HirInfo {
     let (maybe_fn_sig, hir_body) = fn_sig_and_body(tcx, def_id);
 
     let function_source_hash = hash_mir_source(tcx, hir_body);
@@ -70,7 +80,15 @@ fn make_coverage_hir_info(tcx: TyCtxt<'_>, def_id: LocalDefId) -> mir::coverage:
         body_span.shrink_to_lo()
     };
 
-    mir::coverage::HirInfo { function_source_hash, fn_sig_span, body_span }
+    let HirInfoBuilder { branch_coverage_enabled: _, num_block_markers, branch_spans } = builder;
+
+    mir::coverage::HirInfo {
+        function_source_hash,
+        fn_sig_span,
+        body_span,
+        num_block_markers,
+        branch_spans,
+    }
 }
 
 fn fn_sig_and_body(