@@ -13,15 +13,19 @@ use crate::transform::{simplify, MirPass, MirSource};
1313use itertools:: Itertools as _;
1414use rustc:: mir:: * ;
1515use rustc:: ty:: { Ty , TyCtxt } ;
16+ use rustc_index:: vec:: IndexVec ;
1617use rustc_target:: abi:: VariantIdx ;
18+ use std:: iter:: { Enumerate , Peekable } ;
19+ use std:: slice:: Iter ;
1720
1821/// Simplifies arms of form `Variant(x) => Variant(x)` to just a move.
1922///
2023/// This is done by transforming basic blocks where the statements match:
2124///
2225/// ```rust
2326/// _LOCAL_TMP = ((_LOCAL_1 as Variant ).FIELD: TY );
24- /// ((_LOCAL_0 as Variant).FIELD: TY) = move _LOCAL_TMP;
27+ /// _TMP_2 = _LOCAL_TMP;
28+ /// ((_LOCAL_0 as Variant).FIELD: TY) = move _TMP_2;
2529/// discriminant(_LOCAL_0) = VAR_IDX;
2630/// ```
2731///
@@ -32,50 +36,251 @@ use rustc_target::abi::VariantIdx;
3236/// ```
3337pub struct SimplifyArmIdentity ;
3438
39+ #[ derive( Debug ) ]
40+ struct ArmIdentityInfo < ' tcx > {
41+ /// Storage location for the variant's field
42+ local_temp_0 : Local ,
43+ /// Storage location holding the varient being read from
44+ local_1 : Local ,
45+ /// The varient field being read from
46+ vf_s0 : VarField < ' tcx > ,
47+
48+ /// Tracks each assignment to a temporary of the varient's field
49+ field_tmp_assignments : Vec < ( Local , Local ) > ,
50+
51+ /// Storage location holding the variant's field that was read from
52+ local_tmp_s1 : Local ,
53+ /// Storage location holding the enum that we are writing to
54+ local_0 : Local ,
55+ /// The varient field being written to
56+ vf_s1 : VarField < ' tcx > ,
57+
58+ /// Storage location that the discrimentant is being set to
59+ set_discr_local : Local ,
60+ /// The variant being written
61+ set_discr_var_idx : VariantIdx ,
62+
63+ /// Index of the statement that should be overwritten as a move
64+ stmt_to_overwrite : usize ,
65+ /// SourceInfo for the new move
66+ source_info : SourceInfo ,
67+
68+ /// Indexes of matching Storage{Live,Dead} statements encountered.
69+ /// (StorageLive index,, StorageDead index, Local)
70+ storage_stmts : Vec < ( usize , usize , Local ) > ,
71+
72+ /// The statements that should be removed (turned into nops)
73+ stmts_to_remove : Vec < usize > ,
74+ }
75+
76+ fn get_arm_identity_info < ' a , ' tcx > ( stmts : & ' a [ Statement < ' tcx > ] ) -> Option < ArmIdentityInfo < ' tcx > > {
77+ let mut tmp_assigns = Vec :: new ( ) ;
78+ let mut nop_stmts = Vec :: new ( ) ;
79+ let mut storage_stmts = Vec :: new ( ) ;
80+ let mut storage_live_stmts = Vec :: new ( ) ;
81+ let mut storage_dead_stmts = Vec :: new ( ) ;
82+
83+ type StmtIter < ' a , ' tcx > = Peekable < Enumerate < Iter < ' a , Statement < ' tcx > > > > ;
84+
85+ fn is_storage_stmt < ' tcx > ( stmt : & Statement < ' tcx > ) -> bool {
86+ matches ! ( stmt. kind, StatementKind :: StorageLive ( _) | StatementKind :: StorageDead ( _) )
87+ }
88+
89+ fn try_eat_storage_stmts < ' a , ' tcx > (
90+ stmt_iter : & mut StmtIter < ' a , ' tcx > ,
91+ storage_live_stmts : & mut Vec < ( usize , Local ) > ,
92+ storage_dead_stmts : & mut Vec < ( usize , Local ) > ,
93+ ) {
94+ while stmt_iter. peek ( ) . map ( |( _, stmt) | is_storage_stmt ( stmt) ) . unwrap_or ( false ) {
95+ let ( idx, stmt) = stmt_iter. next ( ) . unwrap ( ) ;
96+
97+ if let StatementKind :: StorageLive ( l) = stmt. kind {
98+ storage_live_stmts. push ( ( idx, l) ) ;
99+ } else if let StatementKind :: StorageDead ( l) = stmt. kind {
100+ storage_dead_stmts. push ( ( idx, l) ) ;
101+ }
102+ }
103+ }
104+
105+ fn is_tmp_storage_stmt < ' tcx > ( stmt : & Statement < ' tcx > ) -> bool {
106+ if let StatementKind :: Assign ( box ( place, Rvalue :: Use ( op) ) ) = & stmt. kind {
107+ if let Operand :: Copy ( p) | Operand :: Move ( p) = op {
108+ return place. as_local ( ) . is_some ( ) && p. as_local ( ) . is_some ( ) ;
109+ }
110+ }
111+
112+ false
113+ }
114+
115+ fn try_eat_assign_tmp_stmts < ' a , ' tcx > (
116+ stmt_iter : & mut StmtIter < ' a , ' tcx > ,
117+ tmp_assigns : & mut Vec < ( Local , Local ) > ,
118+ nop_stmts : & mut Vec < usize > ,
119+ ) {
120+ while stmt_iter. peek ( ) . map ( |( _, stmt) | is_tmp_storage_stmt ( stmt) ) . unwrap_or ( false ) {
121+ let ( idx, stmt) = stmt_iter. next ( ) . unwrap ( ) ;
122+
123+ if let StatementKind :: Assign ( box ( place, Rvalue :: Use ( op) ) ) = & stmt. kind {
124+ if let Operand :: Copy ( p) | Operand :: Move ( p) = op {
125+ tmp_assigns. push ( ( place. as_local ( ) . unwrap ( ) , p. as_local ( ) . unwrap ( ) ) ) ;
126+ nop_stmts. push ( idx) ;
127+ }
128+ }
129+ }
130+ }
131+
132+ let mut stmt_iter = stmts. iter ( ) . enumerate ( ) . peekable ( ) ;
133+
134+ try_eat_storage_stmts ( & mut stmt_iter, & mut storage_live_stmts, & mut storage_dead_stmts) ;
135+
136+ let ( starting_stmt, stmt) = stmt_iter. next ( ) ?;
137+ let ( local_tmp_s0, local_1, vf_s0) = match_get_variant_field ( stmt) ?;
138+
139+ try_eat_storage_stmts ( & mut stmt_iter, & mut storage_live_stmts, & mut storage_dead_stmts) ;
140+
141+ try_eat_assign_tmp_stmts ( & mut stmt_iter, & mut tmp_assigns, & mut nop_stmts) ;
142+
143+ try_eat_storage_stmts ( & mut stmt_iter, & mut storage_live_stmts, & mut storage_dead_stmts) ;
144+
145+ try_eat_assign_tmp_stmts ( & mut stmt_iter, & mut tmp_assigns, & mut nop_stmts) ;
146+
147+ let ( idx, stmt) = stmt_iter. next ( ) ?;
148+ let ( local_tmp_s1, local_0, vf_s1) = match_set_variant_field ( stmt) ?;
149+ nop_stmts. push ( idx) ;
150+
151+ let ( idx, stmt) = stmt_iter. next ( ) ?;
152+ let ( set_discr_local, set_discr_var_idx) = match_set_discr ( stmt) ?;
153+ let discr_stmt_source_info = stmt. source_info ;
154+ nop_stmts. push ( idx) ;
155+
156+ try_eat_storage_stmts ( & mut stmt_iter, & mut storage_live_stmts, & mut storage_dead_stmts) ;
157+
158+ for ( live_idx, live_local) in storage_live_stmts {
159+ if let Some ( i) = storage_dead_stmts. iter ( ) . rposition ( |( _, l) | * l == live_local) {
160+ let ( dead_idx, _) = storage_dead_stmts. swap_remove ( i) ;
161+ storage_stmts. push ( ( live_idx, dead_idx, live_local) ) ;
162+ }
163+ }
164+
165+ Some ( ArmIdentityInfo {
166+ local_temp_0 : local_tmp_s0,
167+ local_1,
168+ vf_s0,
169+ field_tmp_assignments : tmp_assigns,
170+ local_tmp_s1,
171+ local_0,
172+ vf_s1,
173+ set_discr_local,
174+ set_discr_var_idx,
175+ stmt_to_overwrite : starting_stmt,
176+ source_info : discr_stmt_source_info,
177+ storage_stmts,
178+ stmts_to_remove : nop_stmts,
179+ } )
180+ }
181+
182+ fn optimization_applies < ' tcx > (
183+ opt_info : & ArmIdentityInfo < ' tcx > ,
184+ local_decls : & IndexVec < Local , LocalDecl < ' tcx > > ,
185+ ) -> bool {
186+ trace ! ( "testing if optimization applies..." ) ;
187+
188+ if opt_info. local_0 == opt_info. local_1 {
189+ trace ! ( "NO: moving into ourselves" ) ;
190+ return false ;
191+ } else if opt_info. vf_s0 != opt_info. vf_s1 {
192+ trace ! ( "NO: the field-and-variant information do not match" ) ;
193+ return false ;
194+ } else if local_decls[ opt_info. local_0 ] . ty != local_decls[ opt_info. local_1 ] . ty {
195+ // FIXME(Centril,oli-obk): possibly relax ot same layout?
196+ trace ! ( "NO: source and target locals have different types" ) ;
197+ return false ;
198+ } else if ( opt_info. local_0 , opt_info. vf_s0 . var_idx )
199+ != ( opt_info. set_discr_local , opt_info. set_discr_var_idx )
200+ {
201+ trace ! ( "NO: the discriminants do not match" ) ;
202+ return false ;
203+ }
204+
205+ // Verify the assigment chain consists of the form b = a; c = b; d = c; etc...
206+ if opt_info. field_tmp_assignments . len ( ) == 0 {
207+ trace ! ( "NO: no assignments found" ) ;
208+ }
209+ let mut last_assigned_to = opt_info. field_tmp_assignments [ 0 ] . 1 ;
210+ let source_local = last_assigned_to;
211+ for ( l, r) in & opt_info. field_tmp_assignments {
212+ if * r != last_assigned_to {
213+ trace ! ( "NO: found unexpected assignment {:?} = {:?}" , l, r) ;
214+ return false ;
215+ }
216+
217+ last_assigned_to = * l;
218+ }
219+
220+ if source_local != opt_info. local_temp_0 {
221+ trace ! (
222+ "NO: start of assignment chain does not match enum variant temp: {:?} != {:?}" ,
223+ source_local,
224+ opt_info. local_temp_0
225+ ) ;
226+ return false ;
227+ } else if last_assigned_to != opt_info. local_tmp_s1 {
228+ trace ! (
229+ "NO: end of assignemnt chain does not match written enum temp: {:?} != {:?}" ,
230+ last_assigned_to,
231+ opt_info. local_tmp_s1
232+ ) ;
233+ return false ;
234+ }
235+
236+ trace ! ( "SUCCESS: optimization applies!" ) ;
237+ return true ;
238+ }
239+
35240impl < ' tcx > MirPass < ' tcx > for SimplifyArmIdentity {
36- fn run_pass ( & self , _: TyCtxt < ' tcx > , _: MirSource < ' tcx > , body : & mut BodyAndCache < ' tcx > ) {
241+ fn run_pass ( & self , _: TyCtxt < ' tcx > , source : MirSource < ' tcx > , body : & mut BodyAndCache < ' tcx > ) {
242+ trace ! ( "running SimplifyArmIdentity on {:?}" , source) ;
37243 let ( basic_blocks, local_decls) = body. basic_blocks_and_local_decls_mut ( ) ;
38244 for bb in basic_blocks {
39- // Need 3 statements:
40- let ( s0, s1, s2) = match & mut * bb. statements {
41- [ s0, s1, s2] => ( s0, s1, s2) ,
42- _ => continue ,
43- } ;
245+ trace ! ( "bb is {:?}" , bb. statements) ;
44246
45- // Pattern match on the form we want:
46- let ( local_tmp_s0, local_1, vf_s0) = match match_get_variant_field ( s0) {
47- None => continue ,
48- Some ( x) => x,
49- } ;
50- let ( local_tmp_s1, local_0, vf_s1) = match match_set_variant_field ( s1) {
51- None => continue ,
52- Some ( x) => x,
53- } ;
54- if local_tmp_s0 != local_tmp_s1
55- // Avoid moving into ourselves.
56- || local_0 == local_1
57- // The field-and-variant information match up.
58- || vf_s0 != vf_s1
59- // Source and target locals have the same type.
60- // FIXME(Centril | oli-obk): possibly relax to same layout?
61- || local_decls[ local_0] . ty != local_decls[ local_1] . ty
62- // We're setting the discriminant of `local_0` to this variant.
63- || Some ( ( local_0, vf_s0. var_idx ) ) != match_set_discr ( s2)
64- {
65- continue ;
66- }
247+ if let Some ( mut opt_info) = get_arm_identity_info ( & bb. statements ) {
248+ trace ! ( "got opt_info = {:#?}" , opt_info) ;
249+ if !optimization_applies ( & opt_info, local_decls) {
250+ debug ! ( "optimization skipped for {:?}" , source) ;
251+ continue ;
252+ }
253+
254+ // Also remove unused Storage{Live,Dead} statements which correspond
255+ // to temps used previously.
256+ for ( left, right) in opt_info. field_tmp_assignments {
257+ for ( live_idx, dead_idx, local) in & opt_info. storage_stmts {
258+ if * local == left || * local == right {
259+ opt_info. stmts_to_remove . push ( * live_idx) ;
260+ opt_info. stmts_to_remove . push ( * dead_idx) ;
261+ }
262+ }
263+ }
67264
68- // Right shape; transform!
69- s0. source_info = s2. source_info ;
70- match & mut s0. kind {
71- StatementKind :: Assign ( box ( place, rvalue) ) => {
72- * place = local_0. into ( ) ;
73- * rvalue = Rvalue :: Use ( Operand :: Move ( local_1. into ( ) ) ) ;
265+ // Right shape; transform!
266+ let stmt = & mut bb. statements [ opt_info. stmt_to_overwrite ] ;
267+ stmt. source_info = opt_info. source_info ;
268+ match & mut stmt. kind {
269+ StatementKind :: Assign ( box ( place, rvalue) ) => {
270+ * place = opt_info. local_0 . into ( ) ;
271+ * rvalue = Rvalue :: Use ( Operand :: Move ( opt_info. local_1 . into ( ) ) ) ;
272+ }
273+ _ => unreachable ! ( ) ,
74274 }
75- _ => unreachable ! ( ) ,
275+
276+ for stmt_idx in opt_info. stmts_to_remove {
277+ bb. statements [ stmt_idx] . make_nop ( ) ;
278+ }
279+
280+ bb. statements . retain ( |stmt| stmt. kind != StatementKind :: Nop ) ;
281+
282+ trace ! ( "block is now {:?}" , bb. statements) ;
76283 }
77- s1. make_nop ( ) ;
78- s2. make_nop ( ) ;
79284 }
80285 }
81286}
@@ -129,7 +334,7 @@ fn match_set_discr<'tcx>(stmt: &Statement<'tcx>) -> Option<(Local, VariantIdx)>
129334 }
130335}
131336
132- #[ derive( PartialEq ) ]
337+ #[ derive( PartialEq , Debug ) ]
133338struct VarField < ' tcx > {
134339 field : Field ,
135340 field_ty : Ty < ' tcx > ,
0 commit comments