|
| 1 | +use std::assert_matches::assert_matches; |
| 2 | +use std::collections::hash_map::Entry; |
| 3 | + |
| 4 | +use rustc_data_structures::fx::FxHashMap; |
| 5 | +use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind}; |
| 6 | +use rustc_middle::mir::{self, BasicBlock, UnOp}; |
| 7 | +use rustc_middle::thir::{ExprId, ExprKind, Thir}; |
| 8 | +use rustc_middle::ty::TyCtxt; |
| 9 | +use rustc_span::def_id::LocalDefId; |
| 10 | + |
| 11 | +use crate::build::Builder; |
| 12 | + |
| 13 | +pub(crate) struct BranchInfoBuilder { |
| 14 | + /// Maps condition expressions to their enclosing `!`, for better instrumentation. |
| 15 | + nots: FxHashMap<ExprId, NotInfo>, |
| 16 | + |
| 17 | + num_block_markers: usize, |
| 18 | + branch_spans: Vec<BranchSpan>, |
| 19 | +} |
| 20 | + |
| 21 | +#[derive(Clone, Copy)] |
| 22 | +struct NotInfo { |
| 23 | + /// When visiting the associated expression as a branch condition, treat this |
| 24 | + /// enclosing `!` as the branch condition instead. |
| 25 | + enclosing_not: ExprId, |
| 26 | + /// True if the associated expression is nested within an odd number of `!` |
| 27 | + /// expressions relative to `enclosing_not` (inclusive of `enclosing_not`). |
| 28 | + is_flipped: bool, |
| 29 | +} |
| 30 | + |
| 31 | +impl BranchInfoBuilder { |
| 32 | + /// Creates a new branch info builder, but only if branch coverage instrumentation |
| 33 | + /// is enabled and `def_id` represents a function that is eligible for coverage. |
| 34 | + pub(crate) fn new_if_enabled(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<Self> { |
| 35 | + if tcx.sess.instrument_coverage_branch() && tcx.is_eligible_for_coverage(def_id) { |
| 36 | + Some(Self { nots: FxHashMap::default(), num_block_markers: 0, branch_spans: vec![] }) |
| 37 | + } else { |
| 38 | + None |
| 39 | + } |
| 40 | + } |
| 41 | + |
| 42 | + /// Unary `!` expressions inside an `if` condition are lowered by lowering |
| 43 | + /// their argument instead, and then reversing the then/else arms of that `if`. |
| 44 | + /// |
| 45 | + /// That's awkward for branch coverage instrumentation, so to work around that |
| 46 | + /// we pre-emptively visit any affected `!` expressions, and record extra |
| 47 | + /// information that [`Builder::visit_coverage_branch_condition`] can use to |
| 48 | + /// synthesize branch instrumentation for the enclosing `!`. |
| 49 | + pub(crate) fn visit_unary_not(&mut self, thir: &Thir<'_>, unary_not: ExprId) { |
| 50 | + assert_matches!(thir[unary_not].kind, ExprKind::Unary { op: UnOp::Not, .. }); |
| 51 | + |
| 52 | + self.visit_with_not_info( |
| 53 | + thir, |
| 54 | + unary_not, |
| 55 | + // Set `is_flipped: false` for the `!` itself, so that its enclosed |
| 56 | + // expression will have `is_flipped: true`. |
| 57 | + NotInfo { enclosing_not: unary_not, is_flipped: false }, |
| 58 | + ); |
| 59 | + } |
| 60 | + |
| 61 | + fn visit_with_not_info(&mut self, thir: &Thir<'_>, expr_id: ExprId, not_info: NotInfo) { |
| 62 | + match self.nots.entry(expr_id) { |
| 63 | + // This expression has already been marked by an enclosing `!`. |
| 64 | + Entry::Occupied(_) => return, |
| 65 | + Entry::Vacant(entry) => entry.insert(not_info), |
| 66 | + }; |
| 67 | + |
| 68 | + match thir[expr_id].kind { |
| 69 | + ExprKind::Unary { op: UnOp::Not, arg } => { |
| 70 | + // Invert the `is_flipped` flag for the contents of this `!`. |
| 71 | + let not_info = NotInfo { is_flipped: !not_info.is_flipped, ..not_info }; |
| 72 | + self.visit_with_not_info(thir, arg, not_info); |
| 73 | + } |
| 74 | + ExprKind::Scope { value, .. } => self.visit_with_not_info(thir, value, not_info), |
| 75 | + ExprKind::Use { source } => self.visit_with_not_info(thir, source, not_info), |
| 76 | + // All other expressions (including `&&` and `||`) don't need any |
| 77 | + // special handling of their contents, so stop visiting. |
| 78 | + _ => {} |
| 79 | + } |
| 80 | + } |
| 81 | + |
| 82 | + fn next_block_marker_id(&mut self) -> BlockMarkerId { |
| 83 | + let id = BlockMarkerId::from_usize(self.num_block_markers); |
| 84 | + self.num_block_markers += 1; |
| 85 | + id |
| 86 | + } |
| 87 | + |
| 88 | + pub(crate) fn into_done(self) -> Option<Box<mir::coverage::BranchInfo>> { |
| 89 | + let Self { nots: _, num_block_markers, branch_spans } = self; |
| 90 | + |
| 91 | + if num_block_markers == 0 { |
| 92 | + assert!(branch_spans.is_empty()); |
| 93 | + return None; |
| 94 | + } |
| 95 | + |
| 96 | + Some(Box::new(mir::coverage::BranchInfo { num_block_markers, branch_spans })) |
| 97 | + } |
| 98 | +} |
| 99 | + |
| 100 | +impl Builder<'_, '_> { |
| 101 | + /// If branch coverage is enabled, inject marker statements into `then_block` |
| 102 | + /// and `else_block`, and record their IDs in the table of branch spans. |
| 103 | + pub(crate) fn visit_coverage_branch_condition( |
| 104 | + &mut self, |
| 105 | + mut expr_id: ExprId, |
| 106 | + mut then_block: BasicBlock, |
| 107 | + mut else_block: BasicBlock, |
| 108 | + ) { |
| 109 | + // Bail out if branch coverage is not enabled for this function. |
| 110 | + let Some(branch_info) = self.coverage_branch_info.as_ref() else { return }; |
| 111 | + |
| 112 | + // If this condition expression is nested within one or more `!` expressions, |
| 113 | + // replace it with the enclosing `!` collected by `visit_unary_not`. |
| 114 | + if let Some(&NotInfo { enclosing_not, is_flipped }) = branch_info.nots.get(&expr_id) { |
| 115 | + expr_id = enclosing_not; |
| 116 | + if is_flipped { |
| 117 | + std::mem::swap(&mut then_block, &mut else_block); |
| 118 | + } |
| 119 | + } |
| 120 | + let source_info = self.source_info(self.thir[expr_id].span); |
| 121 | + |
| 122 | + // Now that we have `source_info`, we can upgrade to a &mut reference. |
| 123 | + let branch_info = self.coverage_branch_info.as_mut().expect("upgrading & to &mut"); |
| 124 | + |
| 125 | + let mut inject_branch_marker = |block: BasicBlock| { |
| 126 | + let id = branch_info.next_block_marker_id(); |
| 127 | + |
| 128 | + let marker_statement = mir::Statement { |
| 129 | + source_info, |
| 130 | + kind: mir::StatementKind::Coverage(Box::new(mir::Coverage { |
| 131 | + kind: CoverageKind::BlockMarker { id }, |
| 132 | + })), |
| 133 | + }; |
| 134 | + self.cfg.push(block, marker_statement); |
| 135 | + |
| 136 | + id |
| 137 | + }; |
| 138 | + |
| 139 | + let true_marker = inject_branch_marker(then_block); |
| 140 | + let false_marker = inject_branch_marker(else_block); |
| 141 | + |
| 142 | + branch_info.branch_spans.push(BranchSpan { |
| 143 | + span: source_info.span, |
| 144 | + true_marker, |
| 145 | + false_marker, |
| 146 | + }); |
| 147 | + } |
| 148 | +} |
0 commit comments