coverage: Record branch information during MIR building

Author: Zalathar

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -1,8 +1,11 @@
-use rustc_middle::mir;
-use rustc_middle::mir::coverage::BranchSpan;
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
+use rustc_middle::mir::{self, BasicBlock, Statement, StatementKind, UnOp};
+use rustc_middle::thir::{ExprId, ExprKind};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::LocalDefId;
 
+use crate::build::Builder;
+
 pub(crate) struct HirBranchInfoBuilder {
     num_block_markers: usize,
     branch_spans: Vec<BranchSpan>,
@@ -17,6 +20,12 @@ impl HirBranchInfoBuilder {
         }
     }
 
+    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 into_done(self) -> Option<Box<mir::coverage::HirBranchInfo>> {
         let Self { num_block_markers, branch_spans } = self;
 
@@ -28,3 +37,72 @@ impl HirBranchInfoBuilder {
         Some(Box::new(mir::coverage::HirBranchInfo { num_block_markers, branch_spans }))
     }
 }
+
+impl<'a, 'tcx> Builder<'a, 'tcx> {
+    /// If branch coverage is enabled, inject marker statements into `then_block`
+    /// and `else_block`, and record their IDs in the table of branch spans.
+    pub(crate) fn coverage_record_branch(
+        &mut self,
+        cond_expr_id: ExprId,
+        enclosing_not_expr_id: Option<ExprId>,
+        then_block: BasicBlock,
+        else_block: BasicBlock,
+    ) {
+        // If branch coverage instrumentation isn't enabled, do nothing.
+        if self.coverage_branch_info.is_none() {
+            return;
+        }
+
+        // If the condition is nested in an odd number of `!` expressions, we
+        // need to reverse the meaning of the then/else block markers, so that
+        // they correspond to the outermost `!` expression being true/false.
+        let swap_arm_markers = {
+            let mut curr = enclosing_not_expr_id.unwrap_or(cond_expr_id);
+            let mut swap_arm_markers = false;
+
+            while curr != cond_expr_id {
+                match self.thir[curr].kind {
+                    ExprKind::Scope { value, .. } => curr = value,
+                    ExprKind::Unary { op: UnOp::Not, arg } => {
+                        swap_arm_markers = !swap_arm_markers;
+                        curr = arg;
+                    }
+                    _ => unreachable!("ensured by `Builder::then_else_break_inner`"),
+                }
+            }
+            swap_arm_markers
+        };
+
+        let cond_source_info =
+            self.source_info(self.thir[enclosing_not_expr_id.unwrap_or(cond_expr_id)].span);
+
+        let branch_info_builder =
+            self.coverage_branch_info.as_mut().expect("confirmed present above");
+
+        let mut inject_branch_marker = |block: BasicBlock| {
+            let id = branch_info_builder.next_block_marker_id();
+
+            let marker_statement = Statement {
+                source_info: cond_source_info,
+                kind: StatementKind::Coverage(Box::new(mir::Coverage {
+                    kind: CoverageKind::BlockMarker { id },
+                })),
+            };
+            self.cfg.push(block, marker_statement);
+
+            id
+        };
+
+        let mut true_marker = inject_branch_marker(then_block);
+        let mut false_marker = inject_branch_marker(else_block);
+        if swap_arm_markers {
+            std::mem::swap(&mut true_marker, &mut false_marker);
+        }
+
+        branch_info_builder.branch_spans.push(BranchSpan {
+            span: cond_source_info.span,
+            true_marker,
+            false_marker,
+        });
+    }
+}

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

@@ -38,9 +38,30 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// expressions will not be declared. This is for if let guards on arms with
     /// an or pattern, where the guard is lowered multiple times.
     pub(crate) fn then_else_break(
+        &mut self,
+        block: BasicBlock,
+        expr_id: ExprId,
+        temp_scope_override: Option<region::Scope>,
+        break_scope: region::Scope,
+        variable_source_info: SourceInfo,
+        declare_bindings: bool,
+    ) -> BlockAnd<()> {
+        self.then_else_break_inner(
+            block,
+            expr_id,
+            None,
+            temp_scope_override,
+            break_scope,
+            variable_source_info,
+            declare_bindings,
+        )
+    }
+
+    fn then_else_break_inner(
         &mut self,
         mut block: BasicBlock,
         expr_id: ExprId,
+        enclosing_not_expr_id: Option<ExprId>,
         temp_scope_override: Option<region::Scope>,
         break_scope: region::Scope,
         variable_source_info: SourceInfo,
@@ -52,18 +73,20 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
         match expr.kind {
             ExprKind::LogicalOp { op: LogicalOp::And, lhs, rhs } => {
-                let lhs_then_block = unpack!(this.then_else_break(
+                let lhs_then_block = unpack!(this.then_else_break_inner(
                     block,
                     lhs,
+                    None,
                     temp_scope_override,
                     break_scope,
                     variable_source_info,
                     declare_bindings,
                 ));
 
-                let rhs_then_block = unpack!(this.then_else_break(
+                let rhs_then_block = unpack!(this.then_else_break_inner(
                     lhs_then_block,
                     rhs,
+                    None,
                     temp_scope_override,
                     break_scope,
                     variable_source_info,
@@ -76,18 +99,20 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let local_scope = this.local_scope();
                 let (lhs_success_block, failure_block) =
                     this.in_if_then_scope(local_scope, expr_span, |this| {
-                        this.then_else_break(
+                        this.then_else_break_inner(
                             block,
                             lhs,
+                            None,
                             temp_scope_override,
                             local_scope,
                             variable_source_info,
                             true,
                         )
                     });
-                let rhs_success_block = unpack!(this.then_else_break(
+                let rhs_success_block = unpack!(this.then_else_break_inner(
                     failure_block,
                     rhs,
+                    None,
                     temp_scope_override,
                     break_scope,
                     variable_source_info,
@@ -108,9 +133,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         if this.tcx.sess.instrument_coverage() {
                             this.cfg.push_coverage_span_marker(block, this.source_info(expr_span));
                         }
-                        this.then_else_break(
+                        this.then_else_break_inner(
                             block,
                             arg,
+                            enclosing_not_expr_id.or(Some(expr_id)),
                             temp_scope_override,
                             local_scope,
                             variable_source_info,
@@ -123,19 +149,21 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             ExprKind::Scope { region_scope, lint_level, value } => {
                 let region_scope = (region_scope, this.source_info(expr_span));
                 this.in_scope(region_scope, lint_level, |this| {
-                    this.then_else_break(
+                    this.then_else_break_inner(
                         block,
                         value,
+                        enclosing_not_expr_id,
                         temp_scope_override,
                         break_scope,
                         variable_source_info,
                         declare_bindings,
                     )
                 })
             }
-            ExprKind::Use { source } => this.then_else_break(
+            ExprKind::Use { source } => this.then_else_break_inner(
                 block,
                 source,
+                None,
                 temp_scope_override,
                 break_scope,
                 variable_source_info,
@@ -161,6 +189,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let else_block = this.cfg.start_new_block();
                 let term = TerminatorKind::if_(operand, then_block, else_block);
 
+                this.coverage_record_branch(expr_id, enclosing_not_expr_id, then_block, else_block);
+
                 let source_info = this.source_info(expr_span);
                 this.cfg.terminate(block, source_info, term);
                 this.break_for_else(else_block, break_scope, source_info);