coverage: Record branch information during MIR building

2 files changed, 131 insertions, 4 deletions

diff --git a/compiler/rustc_mir_build/src/build/coverageinfo.rs b/compiler/rustc_mir_build/src/build/coverageinfo.rs
index 3fdf48f54c9..0b8ec234dda 100644
--- a/compiler/rustc_mir_build/src/build/coverageinfo.rs
+++ b/compiler/rustc_mir_build/src/build/coverageinfo.rs
@@ -1,26 +1,92 @@
-use rustc_middle::mir;
-use rustc_middle::mir::coverage::BranchSpan;
+use std::assert_matches::assert_matches;
+use std::collections::hash_map::Entry;
+
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
+use rustc_middle::mir::{self, BasicBlock, UnOp};
+use rustc_middle::thir::{ExprId, ExprKind, Thir};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::LocalDefId;
 
+use crate::build::Builder;
+
 pub(crate) struct BranchInfoBuilder {
+    /// Maps condition expressions to their enclosing `!`, for better instrumentation.
+    nots: FxHashMap<ExprId, NotInfo>,
+
     num_block_markers: usize,
     branch_spans: Vec<BranchSpan>,
 }
 
+#[derive(Clone, Copy)]
+struct NotInfo {
+    /// When visiting the associated expression as a branch condition, treat this
+    /// enclosing `!` as the branch condition instead.
+    enclosing_not: ExprId,
+    /// True if the associated expression is nested within an odd number of `!`
+    /// expressions relative to `enclosing_not` (inclusive of `enclosing_not`).
+    is_flipped: bool,
+}
+
 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![] })
+            Some(Self { nots: FxHashMap::default(), num_block_markers: 0, branch_spans: vec![] })
         } else {
             None
         }
     }
 
+    /// Unary `!` expressions inside an `if` condition are lowered by lowering
+    /// their argument instead, and then reversing the then/else arms of that `if`.
+    ///
+    /// That's awkward for branch coverage instrumentation, so to work around that
+    /// we pre-emptively visit any affected `!` expressions, and record extra
+    /// information that [`Builder::visit_coverage_branch_condition`] can use to
+    /// synthesize branch instrumentation for the enclosing `!`.
+    pub(crate) fn visit_unary_not(&mut self, thir: &Thir<'_>, unary_not: ExprId) {
+        assert_matches!(thir[unary_not].kind, ExprKind::Unary { op: UnOp::Not, .. });
+
+        self.visit_with_not_info(
+            thir,
+            unary_not,
+            // Set `is_flipped: false` for the `!` itself, so that its enclosed
+            // expression will have `is_flipped: true`.
+            NotInfo { enclosing_not: unary_not, is_flipped: false },
+        );
+    }
+
+    fn visit_with_not_info(&mut self, thir: &Thir<'_>, expr_id: ExprId, not_info: NotInfo) {
+        match self.nots.entry(expr_id) {
+            // This expression has already been marked by an enclosing `!`.
+            Entry::Occupied(_) => return,
+            Entry::Vacant(entry) => entry.insert(not_info),
+        };
+
+        match thir[expr_id].kind {
+            ExprKind::Unary { op: UnOp::Not, arg } => {
+                // Invert the `is_flipped` flag for the contents of this `!`.
+                let not_info = NotInfo { is_flipped: !not_info.is_flipped, ..not_info };
+                self.visit_with_not_info(thir, arg, not_info);
+            }
+            ExprKind::Scope { value, .. } => self.visit_with_not_info(thir, value, not_info),
+            ExprKind::Use { source } => self.visit_with_not_info(thir, source, not_info),
+            // All other expressions (including `&&` and `||`) don't need any
+            // special handling of their contents, so stop visiting.
+            _ => {}
+        }
+    }
+
+    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::BranchInfo>> {
-        let Self { num_block_markers, branch_spans } = self;
+        let Self { nots: _, num_block_markers, branch_spans } = self;
 
         if num_block_markers == 0 {
             assert!(branch_spans.is_empty());
@@ -30,3 +96,53 @@ impl BranchInfoBuilder {
         Some(Box::new(mir::coverage::BranchInfo { num_block_markers, branch_spans }))
     }
 }
+
+impl Builder<'_, '_> {
+    /// 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 visit_coverage_branch_condition(
+        &mut self,
+        mut expr_id: ExprId,
+        mut then_block: BasicBlock,
+        mut else_block: BasicBlock,
+    ) {
+        // Bail out if branch coverage is not enabled for this function.
+        let Some(branch_info) = self.coverage_branch_info.as_ref() else { return };
+
+        // If this condition expression is nested within one or more `!` expressions,
+        // replace it with the enclosing `!` collected by `visit_unary_not`.
+        if let Some(&NotInfo { enclosing_not, is_flipped }) = branch_info.nots.get(&expr_id) {
+            expr_id = enclosing_not;
+            if is_flipped {
+                std::mem::swap(&mut then_block, &mut else_block);
+            }
+        }
+        let source_info = self.source_info(self.thir[expr_id].span);
+
+        // Now that we have `source_info`, we can upgrade to a &mut reference.
+        let branch_info = self.coverage_branch_info.as_mut().expect("upgrading & to &mut");
+
+        let mut inject_branch_marker = |block: BasicBlock| {
+            let id = branch_info.next_block_marker_id();
+
+            let marker_statement = mir::Statement {
+                source_info,
+                kind: mir::StatementKind::Coverage(Box::new(mir::Coverage {
+                    kind: CoverageKind::BlockMarker { id },
+                })),
+            };
+            self.cfg.push(block, marker_statement);
+
+            id
+        };
+
+        let true_marker = inject_branch_marker(then_block);
+        let false_marker = inject_branch_marker(else_block);
+
+        branch_info.branch_spans.push(BranchSpan {
+            span: source_info.span,
+            true_marker,
+            false_marker,
+        });
+    }
+}
diff --git a/compiler/rustc_mir_build/src/build/matches/mod.rs b/compiler/rustc_mir_build/src/build/matches/mod.rs
index 4d5ed65c841..e7808ff880b 100644
--- a/compiler/rustc_mir_build/src/build/matches/mod.rs
+++ b/compiler/rustc_mir_build/src/build/matches/mod.rs
@@ -105,6 +105,13 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 success_block.unit()
             }
             ExprKind::Unary { op: UnOp::Not, arg } => {
+                // Improve branch coverage instrumentation by noting conditions
+                // nested within one or more `!` expressions.
+                // (Skipped if branch coverage is not enabled.)
+                if let Some(branch_info) = this.coverage_branch_info.as_mut() {
+                    branch_info.visit_unary_not(this.thir, expr_id);
+                }
+
                 let local_scope = this.local_scope();
                 let (success_block, failure_block) =
                     this.in_if_then_scope(local_scope, expr_span, |this| {
@@ -149,6 +156,10 @@ 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.
+                // (Does nothing if branch coverage is not enabled.)
+                this.visit_coverage_branch_condition(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);