coverage: Data structures for recording branch info during MIR building

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

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

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;