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

@@ -34,6 +34,7 @@ use std::mem;
 /// to recursive invocations.
 #[derive(Clone, Copy)]
 struct ThenElseArgs {
+    enclosing_not_expr_id: Option<ExprId>,
     /// Used as the temp scope for lowering `expr`. If absent (for match guards),
     /// `self.local_scope()` is used.
     temp_scope_override: Option<region::Scope>,
@@ -61,7 +62,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         self.then_else_break_inner(
             block,
             expr_id,
-            ThenElseArgs { temp_scope_override, variable_source_info, declare_let_bindings },
+            ThenElseArgs {
+                enclosing_not_expr_id: None,
+                temp_scope_override,
+                variable_source_info,
+                declare_let_bindings,
+            },
         )
     }
 
@@ -77,6 +83,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
         match expr.kind {
             ExprKind::LogicalOp { op: LogicalOp::And, lhs, rhs } => {
+                let args = ThenElseArgs { enclosing_not_expr_id: None, ..args };
                 let lhs_then_block = unpack!(this.then_else_break_inner(block, lhs, args));
                 let rhs_then_block = unpack!(this.then_else_break_inner(lhs_then_block, rhs, args));
                 rhs_then_block.unit()
@@ -88,13 +95,21 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         this.then_else_break_inner(
                             block,
                             lhs,
-                            ThenElseArgs { declare_let_bindings: true, ..args },
+                            ThenElseArgs {
+                                enclosing_not_expr_id: None,
+                                declare_let_bindings: true,
+                                ..args
+                            },
                         )
                     });
                 let rhs_success_block = unpack!(this.then_else_break_inner(
                     failure_block,
                     rhs,
-                    ThenElseArgs { declare_let_bindings: true, ..args },
+                    ThenElseArgs {
+                        enclosing_not_expr_id: None,
+                        declare_let_bindings: true,
+                        ..args
+                    },
                 ));
 
                 // Make the LHS and RHS success arms converge to a common block.
@@ -117,7 +132,11 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         this.then_else_break_inner(
                             block,
                             arg,
-                            ThenElseArgs { declare_let_bindings: true, ..args },
+                            ThenElseArgs {
+                                enclosing_not_expr_id: args.enclosing_not_expr_id.or(Some(expr_id)),
+                                declare_let_bindings: true,
+                                ..args
+                            },
                         )
                     });
                 this.break_for_else(success_block, args.variable_source_info);
@@ -150,6 +169,15 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let else_block = this.cfg.start_new_block();
                 let term = TerminatorKind::if_(operand, then_block, else_block);
 
+                // Record branch coverage info for this condition.
+                // This is a no-op if branch coverage is not enabled.
+                this.coverage_record_branch(
+                    expr_id,
+                    args.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, source_info);