@@ -13,6 +13,7 @@ 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 ;
1718
1819/// Simplifies arms of form `Variant(x) => Variant(x)` to just a move.
@@ -21,7 +22,8 @@ use rustc_target::abi::VariantIdx;
2122///
2223/// ```rust
2324/// _LOCAL_TMP = ((_LOCAL_1 as Variant ).FIELD: TY );
24- /// ((_LOCAL_0 as Variant).FIELD: TY) = move _LOCAL_TMP;
25+ /// _TMP_2 = _LOCAL_TMP;
26+ /// ((_LOCAL_0 as Variant).FIELD: TY) = move _TMP_2;
2527/// discriminant(_LOCAL_0) = VAR_IDX;
2628/// ```
2729///
@@ -32,50 +34,213 @@ use rustc_target::abi::VariantIdx;
3234/// ```
3335pub struct SimplifyArmIdentity ;
3436
37+ #[ derive( Debug ) ]
38+ struct ArmIdentityInfo < ' tcx > {
39+ /// Storage location for the variant's field
40+ local_temp_0 : Local ,
41+ /// Storage location holding the varient being read from
42+ local_1 : Local ,
43+ /// The varient field being read from
44+ vf_s0 : VarField < ' tcx > ,
45+
46+ /// Tracks each assignment to a temporary of the varient's field
47+ field_tmp_assignments : Vec < ( Local , Local ) > ,
48+
49+ /// Storage location holding the variant's field that was read from
50+ local_tmp_s1 : Local ,
51+ /// Storage location holding the enum that we are writing to
52+ local_0 : Local ,
53+ /// The varient field being written to
54+ vf_s1 : VarField < ' tcx > ,
55+
56+ /// Storage location that the discrimentant is being set to
57+ set_discr_local : Local ,
58+ /// The variant being written
59+ set_discr_var_idx : VariantIdx ,
60+
61+ /// Index of the statement that should be overwritten as a move
62+ stmt_to_overwrite : usize ,
63+ /// SourceInfo for the new move
64+ source_info : SourceInfo ,
65+
66+ /// Indexes of matching Storage{Live,Dead} statements encountered.
67+ /// (StorageLive index,, StorageDead index, Local)
68+ storage_stmts : Vec < ( usize , usize , Local ) > ,
69+
70+ /// The statements that should be removed (turned into nops)
71+ stmts_to_remove : Vec < usize > ,
72+ }
73+
74+ fn get_arm_identity_info ( stmts : & [ Statement < ' tcx > ] ) -> Option < ArmIdentityInfo < ' tcx > > {
75+ let ( mut local_tmp_s0, mut local_1, mut vf_s0) = ( None , None , None ) ;
76+ let mut tmp_assigns = Vec :: new ( ) ;
77+ let ( mut local_tmp_s1, mut local_0, mut vf_s1) = ( None , None , None ) ;
78+ let ( mut set_discr_local, mut set_discr_var_idx) = ( None , None ) ;
79+ let mut starting_stmt = None ;
80+ let mut discr_stmt = None ;
81+ let mut nop_stmts = Vec :: new ( ) ;
82+ let mut storage_stmts = Vec :: new ( ) ;
83+ let mut storage_live_stmts = Vec :: new ( ) ;
84+ let mut storage_dead_stmts = Vec :: new ( ) ;
85+
86+ for ( stmt_idx, stmt) in stmts. iter ( ) . enumerate ( ) {
87+ if let StatementKind :: StorageLive ( l) = stmt. kind {
88+ storage_live_stmts. push ( ( stmt_idx, l) ) ;
89+ continue ;
90+ } else if let StatementKind :: StorageDead ( l) = stmt. kind {
91+ storage_dead_stmts. push ( ( stmt_idx, l) ) ;
92+ continue ;
93+ }
94+
95+ if local_tmp_s0 == None && local_1 == None && vf_s0 == None {
96+ let result = match_get_variant_field ( stmt) ?;
97+ local_tmp_s0 = Some ( result. 0 ) ;
98+ local_1 = Some ( result. 1 ) ;
99+ vf_s0 = Some ( result. 2 ) ;
100+ starting_stmt = Some ( stmt_idx) ;
101+ } else if let StatementKind :: Assign ( box ( place, Rvalue :: Use ( op) ) ) = & stmt. kind {
102+ if let Some ( local) = place. as_local ( ) {
103+ if let Operand :: Copy ( p) | Operand :: Move ( p) = op {
104+ tmp_assigns. push ( ( local, p. as_local ( ) ?) ) ;
105+ nop_stmts. push ( stmt_idx) ;
106+ } else {
107+ return None ;
108+ }
109+ } else if local_tmp_s1 == None && local_0 == None && vf_s1 == None {
110+ let result = match_set_variant_field ( stmt) ?;
111+ local_tmp_s1 = Some ( result. 0 ) ;
112+ local_0 = Some ( result. 1 ) ;
113+ vf_s1 = Some ( result. 2 ) ;
114+ nop_stmts. push ( stmt_idx) ;
115+ }
116+ } else if set_discr_local == None && set_discr_var_idx == None {
117+ let result = match_set_discr ( stmt) ?;
118+ set_discr_local = Some ( result. 0 ) ;
119+ set_discr_var_idx = Some ( result. 1 ) ;
120+ discr_stmt = Some ( stmt) ;
121+ nop_stmts. push ( stmt_idx) ;
122+ }
123+ }
124+
125+ for ( live_idx, live_local) in storage_live_stmts {
126+ if let Some ( i) = storage_dead_stmts. iter ( ) . rposition ( |( _, l) | * l == live_local) {
127+ let ( dead_idx, _) = storage_dead_stmts. swap_remove ( i) ;
128+ storage_stmts. push ( ( live_idx, dead_idx, live_local) ) ;
129+ }
130+ }
131+
132+ Some ( ArmIdentityInfo {
133+ local_temp_0 : local_tmp_s0?,
134+ local_1 : local_1?,
135+ vf_s0 : vf_s0?,
136+ field_tmp_assignments : tmp_assigns,
137+ local_tmp_s1 : local_tmp_s1?,
138+ local_0 : local_0?,
139+ vf_s1 : vf_s1?,
140+ set_discr_local : set_discr_local?,
141+ set_discr_var_idx : set_discr_var_idx?,
142+ stmt_to_overwrite : starting_stmt?,
143+ source_info : discr_stmt?. source_info ,
144+ storage_stmts : storage_stmts,
145+ stmts_to_remove : nop_stmts,
146+ } )
147+ }
148+
149+ fn optimization_applies < ' tcx > ( opt_info : & ArmIdentityInfo < ' tcx > , local_decls : & IndexVec < Local , LocalDecl < ' tcx > > ) -> bool {
150+ trace ! ( "testing if optimization applies..." ) ;
151+
152+ if opt_info. local_0 == opt_info. local_1 {
153+ trace ! ( "NO: moving into ourselves" ) ;
154+ return false ;
155+ } else if opt_info. vf_s0 != opt_info. vf_s1 {
156+ trace ! ( "NO: the field-and-variant information do not match" ) ;
157+ return false ;
158+ } else if local_decls[ opt_info. local_0 ] . ty != local_decls[ opt_info. local_1 ] . ty {
159+ // FIXME(Centril,oli-obk): possibly relax ot same layout?
160+ trace ! ( "NO: source and target locals have different types" ) ;
161+ return false ;
162+ } else if ( opt_info. local_0 , opt_info. vf_s0 . var_idx ) != ( opt_info. set_discr_local , opt_info. set_discr_var_idx ) {
163+ trace ! ( "NO: the discriminants do not match" ) ;
164+ return false ;
165+ }
166+
167+ // Verify the assigment chain consists of the form b = a; c = b; d = c; etc...
168+ if opt_info. field_tmp_assignments . len ( ) == 0 {
169+ trace ! ( "NO: no assignments found" ) ;
170+ }
171+ let mut last_assigned_to = opt_info. field_tmp_assignments [ 0 ] . 1 ;
172+ let source_local = last_assigned_to;
173+ for ( l, r) in & opt_info. field_tmp_assignments {
174+ if * r != last_assigned_to {
175+ trace ! ( "NO: found unexpected assignment {:?} = {:?}" , l, r) ;
176+ return false ;
177+ }
178+
179+ last_assigned_to = * l;
180+ }
181+
182+ if source_local != opt_info. local_temp_0 {
183+ trace ! ( "NO: start of assignment chain does not match enum variant temp: {:?} != {:?}" , source_local, opt_info. local_temp_0) ;
184+ return false ;
185+ } else if last_assigned_to != opt_info. local_tmp_s1 {
186+ trace ! ( "NO: end of assignemnt chain does not match written enum temp: {:?} != {:?}" , last_assigned_to, opt_info. local_tmp_s1) ;
187+ return false ;
188+ }
189+
190+ trace ! ( "SUCCESS: optimization applies!" ) ;
191+ return true ;
192+ }
193+
35194impl < ' tcx > MirPass < ' tcx > for SimplifyArmIdentity {
36- fn run_pass ( & self , _: TyCtxt < ' tcx > , _: MirSource < ' tcx > , body : & mut BodyAndCache < ' tcx > ) {
195+ fn run_pass ( & self , _: TyCtxt < ' tcx > , source : MirSource < ' tcx > , body : & mut BodyAndCache < ' tcx > ) {
196+ trace ! ( "running SimplifyArmIdentity on {:?}" , source) ;
37197 let ( basic_blocks, local_decls) = body. basic_blocks_and_local_decls_mut ( ) ;
38198 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- } ;
199+ trace ! ( "bb.len() = {:?}" , bb. statements. len( ) ) ;
44200
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- }
201+ if let Some ( mut opt_info) = get_arm_identity_info ( & bb. statements ) {
202+ trace ! ( "got opt_info = {:#?}" , opt_info) ;
203+ if !optimization_applies ( & opt_info, local_decls) {
204+ debug ! ( "skipping simplification!!!!!!!!!!!" ) ;
205+ continue ;
206+ }
207+
208+ trace ! ( "proceeding..." ) ;
209+
210+ //if tcx.sess.opts.debugging_opts.mir_opt_level <= 1 {
211+ // continue;
212+ //}
67213
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 ( ) ) ) ;
214+ // Also remove unused Storage{Live,Dead} statements which correspond
215+ // to temps used previously.
216+ for ( left, right) in opt_info. field_tmp_assignments {
217+ for ( live_idx, dead_idx, local) in & opt_info. storage_stmts {
218+ if * local == left || * local == right {
219+ opt_info. stmts_to_remove . push ( * live_idx) ;
220+ opt_info. stmts_to_remove . push ( * dead_idx) ;
221+ }
222+ }
74223 }
75- _ => unreachable ! ( ) ,
224+
225+ // Right shape; transform!
226+ let stmt = & mut bb. statements [ opt_info. stmt_to_overwrite ] ;
227+ stmt. source_info = opt_info. source_info ;
228+ match & mut stmt. kind {
229+ StatementKind :: Assign ( box ( place, rvalue) ) => {
230+ * place = opt_info. local_0 . into ( ) ;
231+ * rvalue = Rvalue :: Use ( Operand :: Move ( opt_info. local_1 . into ( ) ) ) ;
232+ }
233+ _ => unreachable ! ( ) ,
234+ }
235+
236+ for stmt_idx in opt_info. stmts_to_remove {
237+ bb. statements [ stmt_idx] . make_nop ( ) ;
238+ }
239+
240+ bb. statements . retain ( |stmt| stmt. kind != StatementKind :: Nop ) ;
241+
242+ trace ! ( "block is now {:?}" , bb. statements) ;
76243 }
77- s1. make_nop ( ) ;
78- s2. make_nop ( ) ;
79244 }
80245 }
81246}
@@ -129,7 +294,7 @@ fn match_set_discr<'tcx>(stmt: &Statement<'tcx>) -> Option<(Local, VariantIdx)>
129294 }
130295}
131296
132- #[ derive( PartialEq ) ]
297+ #[ derive( PartialEq , Debug ) ]
133298struct VarField < ' tcx > {
134299 field : Field ,
135300 field_ty : Ty < ' tcx > ,
0 commit comments