@@ -3,12 +3,12 @@ use std::collections::hash_map::Entry;
3
3
4
4
use rustc_data_structures:: fx:: FxHashMap ;
5
5
use rustc_middle:: mir:: coverage:: { BlockMarkerId , BranchSpan , CoverageKind } ;
6
- use rustc_middle:: mir:: { self , BasicBlock , UnOp } ;
6
+ use rustc_middle:: mir:: { self , BasicBlock , SourceInfo , SourceScope , UnOp } ;
7
7
use rustc_middle:: thir:: { ExprId , ExprKind , Thir } ;
8
8
use rustc_middle:: ty:: TyCtxt ;
9
9
use rustc_span:: def_id:: LocalDefId ;
10
10
11
- use crate :: build:: Builder ;
11
+ use crate :: build:: { Builder , CFG } ;
12
12
13
13
pub ( crate ) struct BranchInfoBuilder {
14
14
/// Maps condition expressions to their enclosing `!`, for better instrumentation.
@@ -79,12 +79,55 @@ impl BranchInfoBuilder {
79
79
}
80
80
}
81
81
82
+ fn visit_coverage_branch_condition < ' tcx > (
83
+ & mut self ,
84
+ cfg : & mut CFG < ' tcx > ,
85
+ thir : & Thir < ' tcx > ,
86
+ scope : SourceScope ,
87
+ mut expr_id : ExprId ,
88
+ mut then_block : BasicBlock ,
89
+ mut else_block : BasicBlock ,
90
+ ) {
91
+ // If this condition expression is nested within one or more `!` expressions,
92
+ // replace it with the enclosing `!` collected by `visit_unary_not`.
93
+ if let Some ( & NotInfo { enclosing_not, is_flipped } ) = self . nots . get ( & expr_id) {
94
+ expr_id = enclosing_not;
95
+ if is_flipped {
96
+ std:: mem:: swap ( & mut then_block, & mut else_block) ;
97
+ }
98
+ }
99
+
100
+ let span = thir[ expr_id] . span ;
101
+ let source_info = SourceInfo { span, scope } ;
102
+ let true_marker = self . inject_block_marker ( cfg, source_info, then_block) ;
103
+ let false_marker = self . inject_block_marker ( cfg, source_info, else_block) ;
104
+
105
+ self . branch_spans . push ( BranchSpan { span, true_marker, false_marker } ) ;
106
+ }
107
+
82
108
fn next_block_marker_id ( & mut self ) -> BlockMarkerId {
83
109
let id = BlockMarkerId :: from_usize ( self . num_block_markers ) ;
84
110
self . num_block_markers += 1 ;
85
111
id
86
112
}
87
113
114
+ fn inject_block_marker (
115
+ & mut self ,
116
+ cfg : & mut CFG < ' _ > ,
117
+ source_info : SourceInfo ,
118
+ block : BasicBlock ,
119
+ ) -> BlockMarkerId {
120
+ let id = self . next_block_marker_id ( ) ;
121
+
122
+ let marker_statement = mir:: Statement {
123
+ source_info,
124
+ kind : mir:: StatementKind :: Coverage ( CoverageKind :: BlockMarker { id } ) ,
125
+ } ;
126
+ cfg. push ( block, marker_statement) ;
127
+
128
+ id
129
+ }
130
+
88
131
pub ( crate ) fn into_done ( self ) -> Option < Box < mir:: coverage:: BranchInfo > > {
89
132
let Self { nots : _, num_block_markers, branch_spans } = self ;
90
133
@@ -102,45 +145,20 @@ impl Builder<'_, '_> {
102
145
/// and `else_block`, and record their IDs in the table of branch spans.
103
146
pub ( crate ) fn visit_coverage_branch_condition (
104
147
& mut self ,
105
- mut expr_id : ExprId ,
106
- mut then_block : BasicBlock ,
107
- mut else_block : BasicBlock ,
148
+ expr_id : ExprId ,
149
+ then_block : BasicBlock ,
150
+ else_block : BasicBlock ,
108
151
) {
109
152
// 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 ( CoverageKind :: BlockMarker { id } ) ,
131
- } ;
132
- self . cfg . push ( block, marker_statement) ;
133
-
134
- id
135
- } ;
136
-
137
- let true_marker = inject_branch_marker ( then_block) ;
138
- let false_marker = inject_branch_marker ( else_block) ;
139
-
140
- branch_info. branch_spans . push ( BranchSpan {
141
- span : source_info. span ,
142
- true_marker,
143
- false_marker,
144
- } ) ;
153
+ let Some ( branch_info) = self . coverage_branch_info . as_mut ( ) else { return } ;
154
+
155
+ branch_info. visit_coverage_branch_condition (
156
+ & mut self . cfg ,
157
+ & self . thir ,
158
+ self . source_scope ,
159
+ expr_id,
160
+ then_block,
161
+ else_block,
162
+ ) ;
145
163
}
146
164
}
0 commit comments