@@ -306,49 +306,82 @@ pub struct SimplifyLocals;
306306impl < ' tcx > MirPass < ' tcx > for SimplifyLocals {
307307 fn run_pass ( & self , tcx : TyCtxt < ' tcx > , source : MirSource < ' tcx > , body : & mut BodyAndCache < ' tcx > ) {
308308 trace ! ( "running SimplifyLocals on {:?}" , source) ;
309- let locals = {
309+
310+ // First, we're going to get a count of *actual* uses for every `Local`.
311+ // Take a look at `DeclMarker::visit_local()` to see exactly what is ignored.
312+ let mut used_locals = {
310313 let read_only_cache = read_only ! ( body) ;
311- let mut marker = DeclMarker { locals : BitSet :: new_empty ( body. local_decls . len ( ) ) , body } ;
314+ let mut marker = DeclMarker :: new ( body) ;
312315 marker. visit_body ( & read_only_cache) ;
313- // Return pointer and arguments are always live
314- marker. locals . insert ( RETURN_PLACE ) ;
315- for arg in body. args_iter ( ) {
316- marker. locals . insert ( arg) ;
317- }
318316
319- marker. locals
317+ marker. local_counts
320318 } ;
321319
322- let map = make_local_map ( & mut body. local_decls , locals) ;
323- // Update references to all vars and tmps now
324- LocalUpdater { map, tcx } . visit_body ( body) ;
325- body. local_decls . shrink_to_fit ( ) ;
320+ let arg_count = body. arg_count ;
321+
322+ // Next, we're going to remove any `Local` with zero actual uses. When we remove those
323+ // `Locals`, we're also going to subtract any uses of other `Locals` from the `used_locals`
324+ // count. For example, if we removed `_2 = discriminant(_1)`, then we'll subtract one from
325+ // `use_counts[_1]`. That in turn might make `_1` unused, so we loop until we hit a
326+ // fixedpoint where there are no more unused locals.
327+ loop {
328+ let mut remove_statements = RemoveStatements :: new ( & mut used_locals, arg_count, tcx) ;
329+ remove_statements. visit_body ( body) ;
330+
331+ if !remove_statements. modified {
332+ break ;
333+ }
334+ }
335+
336+ // Finally, we'll actually do the work of shrinking `body.local_decls` and remapping the `Local`s.
337+ let map = make_local_map ( & mut body. local_decls , used_locals, arg_count) ;
338+
339+ // Only bother running the `LocalUpdater` if we actually found locals to remove.
340+ if map. iter ( ) . any ( Option :: is_none) {
341+ // Update references to all vars and tmps now
342+ let mut updater = LocalUpdater { map, tcx } ;
343+ updater. visit_body ( body) ;
344+
345+ body. local_decls . shrink_to_fit ( ) ;
346+ }
326347 }
327348}
328349
329350/// Construct the mapping while swapping out unused stuff out from the `vec`.
330351fn make_local_map < V > (
331- vec : & mut IndexVec < Local , V > ,
332- mask : BitSet < Local > ,
352+ local_decls : & mut IndexVec < Local , V > ,
353+ used_locals : IndexVec < Local , usize > ,
354+ arg_count : usize ,
333355) -> IndexVec < Local , Option < Local > > {
334- let mut map: IndexVec < Local , Option < Local > > = IndexVec :: from_elem ( None , & * vec ) ;
356+ let mut map: IndexVec < Local , Option < Local > > = IndexVec :: from_elem ( None , & * local_decls ) ;
335357 let mut used = Local :: new ( 0 ) ;
336- for alive_index in mask. iter ( ) {
358+ for ( alive_index, count) in used_locals. iter_enumerated ( ) {
359+ // The `RETURN_PLACE` and arguments are always live.
360+ if alive_index. as_usize ( ) > arg_count && * count == 0 {
361+ continue ;
362+ }
363+
337364 map[ alive_index] = Some ( used) ;
338365 if alive_index != used {
339- vec . swap ( alive_index, used) ;
366+ local_decls . swap ( alive_index, used) ;
340367 }
341368 used. increment_by ( 1 ) ;
342369 }
343- vec . truncate ( used. index ( ) ) ;
370+ local_decls . truncate ( used. index ( ) ) ;
344371 map
345372}
346373
347374struct DeclMarker < ' a , ' tcx > {
348- pub locals : BitSet < Local > ,
375+ pub local_counts : IndexVec < Local , usize > ,
349376 pub body : & ' a Body < ' tcx > ,
350377}
351378
379+ impl < ' a , ' tcx > DeclMarker < ' a , ' tcx > {
380+ pub fn new ( body : & ' a Body < ' tcx > ) -> Self {
381+ Self { local_counts : IndexVec :: from_elem ( 0 , & body. local_decls ) , body }
382+ }
383+ }
384+
352385impl < ' a , ' tcx > Visitor < ' tcx > for DeclMarker < ' a , ' tcx > {
353386 fn visit_local ( & mut self , local : & Local , ctx : PlaceContext , location : Location ) {
354387 // Ignore storage markers altogether, they get removed along with their otherwise unused
@@ -368,51 +401,146 @@ impl<'a, 'tcx> Visitor<'tcx> for DeclMarker<'a, 'tcx> {
368401 if location. statement_index != block. statements . len ( ) {
369402 let stmt = & block. statements [ location. statement_index ] ;
370403
404+ fn can_skip_constant ( c : & ty:: Const < ' tcx > ) -> bool {
405+ // Keep assignments from unevaluated constants around, since the
406+ // evaluation may report errors, even if the use of the constant
407+ // is dead code.
408+ !matches ! ( c. val, ty:: ConstKind :: Unevaluated ( ..) )
409+ }
410+
411+ fn can_skip_operand ( o : & Operand < ' _ > ) -> bool {
412+ match o {
413+ Operand :: Copy ( _) | Operand :: Move ( _) => true ,
414+ Operand :: Constant ( c) => can_skip_constant ( c. literal ) ,
415+ }
416+ }
417+
371418 if let StatementKind :: Assign ( box ( dest, rvalue) ) = & stmt. kind {
372419 if !dest. is_indirect ( ) && dest. local == * local {
373- if let Rvalue :: Use ( Operand :: Constant ( c) ) = rvalue {
374- match c. literal . val {
375- // Keep assignments from unevaluated constants around, since the
376- // evaluation may report errors, even if the use of the constant
377- // is dead code.
378- ty:: ConstKind :: Unevaluated ( ..) => { }
379- _ => {
380- trace ! ( "skipping store of const value {:?} to {:?}" , c, dest) ;
381- return ;
382- }
420+ let can_skip = match rvalue {
421+ Rvalue :: Use ( op) => can_skip_operand ( op) ,
422+ Rvalue :: Discriminant ( _) => true ,
423+ Rvalue :: BinaryOp ( _, l, r) | Rvalue :: CheckedBinaryOp ( _, l, r) => {
424+ can_skip_operand ( l) && can_skip_operand ( r)
383425 }
384- } else if let Rvalue :: Discriminant ( d) = rvalue {
385- trace ! ( "skipping store of discriminant value {:?} to {:?}" , d, dest) ;
426+ Rvalue :: Repeat ( op, c) => can_skip_operand ( op) && can_skip_constant ( c) ,
427+ Rvalue :: AddressOf ( _, _) => true ,
428+ Rvalue :: Len ( _) => true ,
429+ Rvalue :: UnaryOp ( _, op) => can_skip_operand ( op) ,
430+ Rvalue :: Aggregate ( _, operands) => operands. iter ( ) . all ( can_skip_operand) ,
431+
432+ _ => false ,
433+ } ;
434+
435+ if can_skip {
436+ trace ! ( "skipping store of {:?} to {:?}" , rvalue, dest) ;
386437 return ;
387438 }
388439 }
389440 }
390441 }
391442 }
392443
393- self . locals . insert ( * local) ;
444+ self . local_counts [ * local] += 1 ;
394445 }
395446}
396447
397- struct LocalUpdater < ' tcx > {
398- map : IndexVec < Local , Option < Local > > ,
448+ struct StatementDeclMarker < ' a , ' tcx > {
449+ used_locals : & ' a mut IndexVec < Local , usize > ,
450+ statement : & ' a Statement < ' tcx > ,
451+ }
452+
453+ impl < ' a , ' tcx > StatementDeclMarker < ' a , ' tcx > {
454+ pub fn new (
455+ used_locals : & ' a mut IndexVec < Local , usize > ,
456+ statement : & ' a Statement < ' tcx > ,
457+ ) -> Self {
458+ Self { used_locals, statement }
459+ }
460+ }
461+
462+ impl < ' a , ' tcx > Visitor < ' tcx > for StatementDeclMarker < ' a , ' tcx > {
463+ fn visit_local ( & mut self , local : & Local , context : PlaceContext , _location : Location ) {
464+ // Skip the lvalue for assignments
465+ if let StatementKind :: Assign ( box ( p, _) ) = self . statement . kind {
466+ if p. local == * local && context. is_place_assignment ( ) {
467+ return ;
468+ }
469+ }
470+
471+ let use_count = & mut self . used_locals [ * local] ;
472+ // If this is the local we're removing...
473+ if * use_count != 0 {
474+ * use_count -= 1 ;
475+ }
476+ }
477+ }
478+
479+ struct RemoveStatements < ' a , ' tcx > {
480+ used_locals : & ' a mut IndexVec < Local , usize > ,
481+ arg_count : usize ,
399482 tcx : TyCtxt < ' tcx > ,
483+ modified : bool ,
400484}
401485
402- impl < ' tcx > MutVisitor < ' tcx > for LocalUpdater < ' tcx > {
486+ impl < ' a , ' tcx > RemoveStatements < ' a , ' tcx > {
487+ fn new (
488+ used_locals : & ' a mut IndexVec < Local , usize > ,
489+ arg_count : usize ,
490+ tcx : TyCtxt < ' tcx > ,
491+ ) -> Self {
492+ Self { used_locals, arg_count, tcx, modified : false }
493+ }
494+
495+ fn keep_local ( & self , l : Local ) -> bool {
496+ trace ! ( "keep_local({:?}): count: {:?}" , l, self . used_locals[ l] ) ;
497+ l. as_usize ( ) <= self . arg_count || self . used_locals [ l] != 0
498+ }
499+ }
500+
501+ impl < ' a , ' tcx > MutVisitor < ' tcx > for RemoveStatements < ' a , ' tcx > {
403502 fn tcx ( & self ) -> TyCtxt < ' tcx > {
404503 self . tcx
405504 }
406505
407506 fn visit_basic_block_data ( & mut self , block : BasicBlock , data : & mut BasicBlockData < ' tcx > ) {
408507 // Remove unnecessary StorageLive and StorageDead annotations.
409- data. statements . retain ( |stmt| match & stmt. kind {
410- StatementKind :: StorageLive ( l) | StatementKind :: StorageDead ( l) => self . map [ * l] . is_some ( ) ,
411- StatementKind :: Assign ( box ( place, _) ) => self . map [ place. local ] . is_some ( ) ,
412- _ => true ,
508+ let mut i = 0usize ;
509+ data. statements . retain ( |stmt| {
510+ let keep = match & stmt. kind {
511+ StatementKind :: StorageLive ( l) | StatementKind :: StorageDead ( l) => {
512+ self . keep_local ( * l)
513+ }
514+ StatementKind :: Assign ( box ( place, _) ) => self . keep_local ( place. local ) ,
515+ _ => true ,
516+ } ;
517+
518+ if !keep {
519+ trace ! ( "removing statement {:?}" , stmt) ;
520+ self . modified = true ;
521+
522+ let mut visitor = StatementDeclMarker :: new ( self . used_locals , stmt) ;
523+ visitor. visit_statement ( stmt, Location { block, statement_index : i } ) ;
524+ }
525+
526+ i += 1 ;
527+
528+ keep
413529 } ) ;
530+
414531 self . super_basic_block_data ( block, data) ;
415532 }
533+ }
534+
535+ struct LocalUpdater < ' tcx > {
536+ map : IndexVec < Local , Option < Local > > ,
537+ tcx : TyCtxt < ' tcx > ,
538+ }
539+
540+ impl < ' tcx > MutVisitor < ' tcx > for LocalUpdater < ' tcx > {
541+ fn tcx ( & self ) -> TyCtxt < ' tcx > {
542+ self . tcx
543+ }
416544
417545 fn visit_local ( & mut self , l : & mut Local , _: PlaceContext , _: Location ) {
418546 * l = self . map [ * l] . unwrap ( ) ;
0 commit comments