@@ -3,12 +3,12 @@ use std::collections::hash_map::Entry;
33
44use rustc_data_structures:: fx:: FxHashMap ;
55use 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 } ;
77use rustc_middle:: thir:: { ExprId , ExprKind , Thir } ;
88use rustc_middle:: ty:: TyCtxt ;
99use rustc_span:: def_id:: LocalDefId ;
1010
11- use crate :: build:: Builder ;
11+ use crate :: build:: { Builder , CFG } ;
1212
1313pub ( crate ) struct BranchInfoBuilder {
1414 /// Maps condition expressions to their enclosing `!`, for better instrumentation.
@@ -79,12 +79,55 @@ impl BranchInfoBuilder {
7979 }
8080 }
8181
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+
82108 fn next_block_marker_id ( & mut self ) -> BlockMarkerId {
83109 let id = BlockMarkerId :: from_usize ( self . num_block_markers ) ;
84110 self . num_block_markers += 1 ;
85111 id
86112 }
87113
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+
88131 pub ( crate ) fn into_done ( self ) -> Option < Box < mir:: coverage:: BranchInfo > > {
89132 let Self { nots : _, num_block_markers, branch_spans } = self ;
90133
@@ -102,45 +145,20 @@ impl Builder<'_, '_> {
102145 /// and `else_block`, and record their IDs in the table of branch spans.
103146 pub ( crate ) fn visit_coverage_branch_condition (
104147 & 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 ,
108151 ) {
109152 // 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+ ) ;
145163 }
146164}
0 commit comments