@@ -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