@@ -2,6 +2,7 @@ use crate::MirPass;
22use rustc_data_structures:: fx:: { FxIndexMap , IndexEntry } ;
33use rustc_index:: bit_set:: BitSet ;
44use rustc_index:: vec:: IndexVec ;
5+ use rustc_middle:: mir:: patch:: MirPatch ;
56use rustc_middle:: mir:: visit:: * ;
67use rustc_middle:: mir:: * ;
78use rustc_middle:: ty:: TyCtxt ;
@@ -13,7 +14,9 @@ impl<'tcx> MirPass<'tcx> for ScalarReplacementOfAggregates {
1314 sess. mir_opt_level ( ) >= 3
1415 }
1516
17+ #[ instrument( level = "debug" , skip( self , tcx, body) ) ]
1618 fn run_pass ( & self , tcx : TyCtxt < ' tcx > , body : & mut Body < ' tcx > ) {
19+ debug ! ( def_id = ?body. source. def_id( ) ) ;
1720 let escaping = escaping_locals ( & * body) ;
1821 debug ! ( ?escaping) ;
1922 let replacements = compute_flattening ( tcx, body, escaping) ;
@@ -69,15 +72,28 @@ fn escaping_locals(body: &Body<'_>) -> BitSet<Local> {
6972 self . super_rvalue ( rvalue, location)
7073 }
7174
75+ fn visit_assign (
76+ & mut self ,
77+ lvalue : & Place < ' tcx > ,
78+ rvalue : & Rvalue < ' tcx > ,
79+ location : Location ,
80+ ) {
81+ if lvalue. as_local ( ) . is_some ( ) && let Rvalue :: Aggregate ( ..) = rvalue {
82+ // Aggregate assignments are expanded in run_pass.
83+ self . visit_rvalue ( rvalue, location) ;
84+ return ;
85+ }
86+ self . super_assign ( lvalue, rvalue, location)
87+ }
88+
7289 fn visit_statement ( & mut self , statement : & Statement < ' tcx > , location : Location ) {
73- if let StatementKind :: StorageLive ( ..)
74- | StatementKind :: StorageDead ( ..)
75- | StatementKind :: Deinit ( ..) = statement. kind
76- {
90+ match statement. kind {
7791 // Storage statements are expanded in run_pass.
78- return ;
92+ StatementKind :: StorageLive ( ..)
93+ | StatementKind :: StorageDead ( ..)
94+ | StatementKind :: Deinit ( ..) => return ,
95+ _ => self . super_statement ( statement, location) ,
7996 }
80- self . super_statement ( statement, location)
8197 }
8298
8399 fn visit_terminator ( & mut self , terminator : & Terminator < ' tcx > , location : Location ) {
@@ -192,6 +208,7 @@ fn replace_flattened_locals<'tcx>(
192208 replacements,
193209 all_dead_locals,
194210 fragments,
211+ patch : MirPatch :: new ( body) ,
195212 } ;
196213 for ( bb, data) in body. basic_blocks . as_mut_preserves_cfg ( ) . iter_enumerated_mut ( ) {
197214 visitor. visit_basic_block_data ( bb, data) ;
@@ -205,6 +222,7 @@ fn replace_flattened_locals<'tcx>(
205222 for var_debug_info in & mut body. var_debug_info {
206223 visitor. visit_var_debug_info ( var_debug_info) ;
207224 }
225+ visitor. patch . apply ( body) ;
208226}
209227
210228struct ReplacementVisitor < ' tcx , ' ll > {
@@ -218,6 +236,7 @@ struct ReplacementVisitor<'tcx, 'll> {
218236 /// Pre-computed list of all "new" locals for each "old" local. This is used to expand storage
219237 /// and deinit statement and debuginfo.
220238 fragments : IndexVec < Local , Vec < ( & ' tcx [ PlaceElem < ' tcx > ] , Local ) > > ,
239+ patch : MirPatch < ' tcx > ,
221240}
222241
223242impl < ' tcx , ' ll > ReplacementVisitor < ' tcx , ' ll > {
@@ -255,12 +274,63 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
255274 }
256275
257276 fn visit_statement ( & mut self , statement : & mut Statement < ' tcx > , location : Location ) {
258- if let StatementKind :: StorageLive ( ..)
259- | StatementKind :: StorageDead ( ..)
260- | StatementKind :: Deinit ( ..) = statement. kind
261- {
262- // Storage statements are expanded in run_pass.
263- return ;
277+ match statement. kind {
278+ StatementKind :: StorageLive ( l) => {
279+ if self . all_dead_locals . contains ( l) {
280+ let final_locals = & self . fragments [ l] ;
281+ for & ( _, fl) in final_locals {
282+ self . patch . add_statement ( location, StatementKind :: StorageLive ( fl) ) ;
283+ }
284+ statement. make_nop ( ) ;
285+ }
286+ return ;
287+ }
288+ StatementKind :: StorageDead ( l) => {
289+ if self . all_dead_locals . contains ( l) {
290+ let final_locals = & self . fragments [ l] ;
291+ for & ( _, fl) in final_locals {
292+ self . patch . add_statement ( location, StatementKind :: StorageDead ( fl) ) ;
293+ }
294+ statement. make_nop ( ) ;
295+ }
296+ return ;
297+ }
298+ StatementKind :: Deinit ( box ref place) => {
299+ if let Some ( local) = place. as_local ( )
300+ && self . all_dead_locals . contains ( local)
301+ {
302+ let final_locals = & self . fragments [ local] ;
303+ for & ( _, fl) in final_locals {
304+ self . patch . add_statement (
305+ location,
306+ StatementKind :: Deinit ( Box :: new ( fl. into ( ) ) ) ,
307+ ) ;
308+ }
309+ statement. make_nop ( ) ;
310+ return ;
311+ }
312+ }
313+
314+ StatementKind :: Assign ( box ( ref place, Rvalue :: Aggregate ( _, ref operands) ) ) => {
315+ if let Some ( local) = place. as_local ( )
316+ && self . all_dead_locals . contains ( local)
317+ {
318+ let final_locals = & self . fragments [ local] ;
319+ for & ( projection, fl) in final_locals {
320+ let & [ PlaceElem :: Field ( index, _) ] = projection else { bug ! ( ) } ;
321+ let index = index. as_usize ( ) ;
322+ let rvalue = Rvalue :: Use ( operands[ index] . clone ( ) ) ;
323+ self . patch . add_statement (
324+ location,
325+ StatementKind :: Assign ( Box :: new ( ( fl. into ( ) , rvalue) ) ) ,
326+ ) ;
327+ }
328+ statement. make_nop ( ) ;
329+ return ;
330+ }
331+ }
332+
333+ _ => { }
264334 }
265335 self . super_statement ( statement, location)
266336 }
@@ -309,39 +379,6 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
309379 }
310380 }
311381
312- fn visit_basic_block_data ( & mut self , bb : BasicBlock , bbdata : & mut BasicBlockData < ' tcx > ) {
313- self . super_basic_block_data ( bb, bbdata) ;
314-
315- #[ derive( Debug ) ]
316- enum Stmt {
317- StorageLive ,
318- StorageDead ,
319- Deinit ,
320- }
321-
322- bbdata. expand_statements ( |stmt| {
323- let source_info = stmt. source_info ;
324- let ( stmt, origin_local) = match & stmt. kind {
325- StatementKind :: StorageLive ( l) => ( Stmt :: StorageLive , * l) ,
326- StatementKind :: StorageDead ( l) => ( Stmt :: StorageDead , * l) ,
327- StatementKind :: Deinit ( p) if let Some ( l) = p. as_local ( ) => ( Stmt :: Deinit , l) ,
328- _ => return None ,
329- } ;
330- if !self . all_dead_locals . contains ( origin_local) {
331- return None ;
332- }
333- let final_locals = self . fragments . get ( origin_local) ?;
334- Some ( final_locals. iter ( ) . map ( move |& ( _, l) | {
335- let kind = match stmt {
336- Stmt :: StorageLive => StatementKind :: StorageLive ( l) ,
337- Stmt :: StorageDead => StatementKind :: StorageDead ( l) ,
338- Stmt :: Deinit => StatementKind :: Deinit ( Box :: new ( l. into ( ) ) ) ,
339- } ;
340- Statement { source_info, kind }
341- } ) )
342- } ) ;
343- }
344-
345382 fn visit_local ( & mut self , local : & mut Local , _: PlaceContext , _: Location ) {
346383 assert ! ( !self . all_dead_locals. contains( * local) ) ;
347384 }
0 commit comments