coverage: Record branch info during MIR building

Author: Zalathar

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -1,11 +1,13 @@
 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::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
+use rustc_middle::mir::{self, BasicBlock, SourceInfo, Statement, StatementKind};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::LocalDefId;
 use rustc_span::{ExpnKind, Span};
 
+use crate::build::Builder;
+
 pub(crate) struct HirInfoBuilder {
     pub(crate) branch_coverage_enabled: bool,
     pub(crate) num_block_markers: usize,
@@ -25,6 +27,12 @@ impl HirInfoBuilder {
         }
     }
 
+    fn next_block_marker_id(&mut self) -> BlockMarkerId {
+        let id = BlockMarkerId::from_usize(self.num_block_markers);
+        self.num_block_markers += 1;
+        id
+    }
+
     pub(crate) fn finish(self, tcx: TyCtxt<'_>, def_id: LocalDefId) -> Box<mir::coverage::HirInfo> {
         Box::new(make_coverage_hir_info(tcx, def_id, self))
     }
@@ -138,3 +146,38 @@ fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx
         .to_smaller_hash()
         .as_u64()
 }
+
+impl<'a, 'tcx> Builder<'a, 'tcx> {
+    pub(crate) fn coverage_add_branch(
+        &mut self,
+        cond_source_info: SourceInfo,
+        then_blk: BasicBlock,
+        else_blk: BasicBlock,
+    ) {
+        let hir_info_builder = self.coverage.as_mut().unwrap();
+        assert!(hir_info_builder.branch_coverage_enabled);
+
+        let mut inject_branch_marker = |bb: BasicBlock| {
+            let id = hir_info_builder.next_block_marker_id();
+
+            let statement = Statement {
+                source_info: cond_source_info,
+                kind: StatementKind::Coverage(Box::new(mir::Coverage {
+                    kind: CoverageKind::BlockMarker { id },
+                })),
+            };
+
+            self.cfg.push(bb, statement);
+            id
+        };
+
+        let true_marker = inject_branch_marker(then_blk);
+        let false_marker = inject_branch_marker(else_blk);
+
+        hir_info_builder.branch_spans.push(BranchSpan {
+            span: cond_source_info.span,
+            true_marker,
+            false_marker,
+        });
+    }
+}

compiler/rustc_mir_build/src/build/matches/mod.rs

@@ -141,6 +141,13 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let term = TerminatorKind::if_(operand, then_block, else_block);
 
                 let source_info = this.source_info(expr_span);
+
+                if let Some(coverage) = &this.coverage
+                    && coverage.branch_coverage_enabled
+                {
+                    this.coverage_add_branch(source_info, then_block, else_block);
+                }
+
                 this.cfg.terminate(block, source_info, term);
                 this.break_for_else(else_block, break_scope, source_info);