Skip to content

Commit 7fb5187

Browse files
borsbors
committed
Auto merge of #70755 - wesleywiser:simplify_locals_2_electric_boogaloo, r=oli-obk
[mir-opt] Run SimplifyLocals to a fixedpoint and handle most rvalues Follow up to review feedback left on #70595.
2 parents 4e4d49d + 9666d31 commit 7fb5187

File tree

7 files changed

+319
-154
lines changed

7 files changed

+319
-154
lines changed

‎src/librustc_mir/transform/simplify.rs‎

Lines changed: 167 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -306,49 +306,82 @@ pub struct SimplifyLocals;
306306
impl<'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`.
330351
fn 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

347374
struct 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+
352385
impl<'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();

‎src/test/mir-opt/const_allocation3/32bit/rustc.main.ConstProp.after.mir‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,30 +30,30 @@ fn main() -> () {
3030
}
3131

3232
alloc0 (static: FOO, size: 4, align: 4) {
33-
╾alloc10+0╼ │ ╾──╼
33+
╾alloc9+0─╼ │ ╾──╼
3434
}
3535

36-
alloc10 (size: 168, align: 1) {
36+
alloc9 (size: 168, align: 1) {
3737
0x00 │ ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab │ ................
38-
0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾alloc5+0─╼ │ ............╾──╼
38+
0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾alloc4+0─╼ │ ............╾──╼
3939
0x20 │ 01 ef cd ab 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4040
0x30 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4141
0x40 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4242
0x50 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4343
0x60 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4444
0x70 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
45-
0x80 │ 00 00 00 00 00 00 00 00 00 00 ╾alloc7+0─╼ 00 00 │ ..........╾──╼..
46-
0x90 │ ╾alloc8+99╼ 00 00 00 00 00 00 00 00 00 00 00 00 │ ╾──╼............
45+
0x80 │ 00 00 00 00 00 00 00 00 00 00 ╾alloc6+0─╼ 00 00 │ ..........╾──╼..
46+
0x90 │ ╾alloc7+99╼ 00 00 00 00 00 00 00 00 00 00 00 00 │ ╾──╼............
4747
0xa0 │ 00 00 00 00 00 00 00 00 │ ........
4848
}
4949

50-
alloc5 (size: 4, align: 4) {
50+
alloc4 (size: 4, align: 4) {
5151
2a 00 00 00 │ *...
5252
}
5353

54-
alloc7 (fn: main)
54+
alloc6 (fn: main)
5555

56-
alloc8 (size: 100, align: 1) {
56+
alloc7 (size: 100, align: 1) {
5757
0x00 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
5858
0x10 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
5959
0x20 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................

‎src/test/mir-opt/const_allocation3/64bit/rustc.main.ConstProp.after.mir‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,31 +30,31 @@ fn main() -> () {
3030
}
3131

3232
alloc0 (static: FOO, size: 8, align: 8) {
33-
╾──────alloc10+0──────╼ │ ╾──────╼
33+
╾──────alloc9+0───────╼ │ ╾──────╼
3434
}
3535

36-
alloc10 (size: 180, align: 1) {
36+
alloc9 (size: 180, align: 1) {
3737
0x00 │ ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab │ ................
38-
0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾─alloc5+0─ │ ............╾───
38+
0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾─alloc4+0─ │ ............╾───
3939
0x20 │ ──────────╼ 01 ef cd ab 00 00 00 00 00 00 00 00 │ ───╼............
4040
0x30 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4141
0x40 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4242
0x50 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4343
0x60 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4444
0x70 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4545
0x80 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 ╾──── │ ..............╾─
46-
0x90 │ ────alloc7+0────╼ 00 00 ╾──────alloc8+99──────╼ │ ─────╼..╾──────╼
46+
0x90 │ ────alloc6+0────╼ 00 00 ╾──────alloc7+99──────╼ │ ─────╼..╾──────╼
4747
0xa0 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
4848
0xb0 │ 00 00 00 00 │ ....
4949
}
5050

51-
alloc5 (size: 4, align: 4) {
51+
alloc4 (size: 4, align: 4) {
5252
2a 00 00 00 │ *...
5353
}
5454

55-
alloc7 (fn: main)
55+
alloc6 (fn: main)
5656

57-
alloc8 (size: 100, align: 1) {
57+
alloc7 (size: 100, align: 1) {
5858
0x00 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
5959
0x10 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
6060
0x20 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................

0 commit comments

Comments
 (0)