[mir-opt] Run SimplifyLocals to a fixedpoint and handle most rvalues

src/librustc_mir/transform/simplify.rs

@@ -306,49 +306,74 @@ pub struct SimplifyLocals;
 impl<'tcx> MirPass<'tcx> for SimplifyLocals {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut BodyAndCache<'tcx>) {
         trace!("running SimplifyLocals on {:?}", source);
-        let locals = {
+
+        let mut used_locals = {
             let read_only_cache = read_only!(body);
-            let mut marker = DeclMarker { locals: BitSet::new_empty(body.local_decls.len()), body };
+            let mut marker = DeclMarker::new(body);
             marker.visit_body(&read_only_cache);
-            // Return pointer and arguments are always live
-            marker.locals.insert(RETURN_PLACE);
-            for arg in body.args_iter() {
-                marker.locals.insert(arg);
-            }
 
-            marker.locals
+            marker.local_counts
         };
 
-        let map = make_local_map(&mut body.local_decls, locals);
-        // Update references to all vars and tmps now
-        LocalUpdater { map, tcx }.visit_body(body);
-        body.local_decls.shrink_to_fit();
+        let arg_count = body.arg_count;
+
+        loop {
+            let mut remove_statements = RemoveStatements::new(&mut used_locals, arg_count, tcx);
+            remove_statements.visit_body(body);
+
+            if !remove_statements.modified {
+                break;
+            }
+        }
+
+        let map = make_local_map(&mut body.local_decls, used_locals, arg_count);
+
+        // Only bother running the `LocalUpdater` if we actually found locals to remove.
+        if map.iter().any(Option::is_none) {
+            // Update references to all vars and tmps now
+            let mut updater = LocalUpdater { map, tcx };
+            updater.visit_body(body);
+
+            body.local_decls.shrink_to_fit();
+        }
     }
 }
 
 /// Construct the mapping while swapping out unused stuff out from the `vec`.
 fn make_local_map<V>(
-    vec: &mut IndexVec<Local, V>,
-    mask: BitSet<Local>,
+    local_decls: &mut IndexVec<Local, V>,
+    used_locals: IndexVec<Local, usize>,
+    arg_count: usize,
 ) -> IndexVec<Local, Option<Local>> {
-    let mut map: IndexVec<Local, Option<Local>> = IndexVec::from_elem(None, &*vec);
+    let mut map: IndexVec<Local, Option<Local>> = IndexVec::from_elem(None, &*local_decls);
     let mut used = Local::new(0);
-    for alive_index in mask.iter() {
+    for (alive_index, count) in used_locals.iter_enumerated() {
+        // The `RETURN_PLACE` and arguments are always live.
+        if alive_index.as_usize() > arg_count && *count == 0 {
+            continue;
+        }
+
         map[alive_index] = Some(used);
         if alive_index != used {
-            vec.swap(alive_index, used);
+            local_decls.swap(alive_index, used);
         }
         used.increment_by(1);
     }
-    vec.truncate(used.index());
+    local_decls.truncate(used.index());
     map
 }
 
 struct DeclMarker<'a, 'tcx> {
-    pub locals: BitSet<Local>,
+    pub local_counts: IndexVec<Local, usize>,
     pub body: &'a Body<'tcx>,
 }
 
+impl<'a, 'tcx> DeclMarker<'a, 'tcx> {
+    pub fn new(body: &'a Body<'tcx>) -> Self {
+        Self { local_counts: IndexVec::from_elem(0, &body.local_decls), body }
+    }
+}
+
 impl<'a, 'tcx> Visitor<'tcx> for DeclMarker<'a, 'tcx> {
     fn visit_local(&mut self, local: &Local, ctx: PlaceContext, location: Location) {
         // Ignore storage markers altogether, they get removed along with their otherwise unused
@@ -368,51 +393,148 @@ impl<'a, 'tcx> Visitor<'tcx> for DeclMarker<'a, 'tcx> {
             if location.statement_index != block.statements.len() {
                 let stmt = &block.statements[location.statement_index];
 
+                fn can_skip_constant(c: &ty::Const<'tcx>) -> bool {
+                    // Keep assignments from unevaluated constants around, since the
+                    // evaluation may report errors, even if the use of the constant
+                    // is dead code.
+                    !matches!(c.val, ty::ConstKind::Unevaluated(..))
+                }
+
+                fn can_skip_operand(o: &Operand<'_>) -> bool {
+                    match o {
+                        Operand::Copy(p) | Operand::Move(p) => !p.is_indirect(),
+                        Operand::Constant(c) => can_skip_constant(c.literal),
+                    }
+                }
+
                 if let StatementKind::Assign(box (dest, rvalue)) = &stmt.kind {
                     if !dest.is_indirect() && dest.local == *local {
-                        if let Rvalue::Use(Operand::Constant(c)) = rvalue {
-                            match c.literal.val {
-                                // Keep assignments from unevaluated constants around, since the
-                                // evaluation may report errors, even if the use of the constant
-                                // is dead code.
-                                ty::ConstKind::Unevaluated(..) => {}
-                                _ => {
-                                    trace!("skipping store of const value {:?} to {:?}", c, dest);
-                                    return;
-                                }
+                        let can_skip = match rvalue {
+                            Rvalue::Use(op) => can_skip_operand(op),
+                            Rvalue::Discriminant(_) => true,
+                            Rvalue::BinaryOp(_, l, r) | Rvalue::CheckedBinaryOp(_, l, r) => {
+                                can_skip_operand(l) && can_skip_operand(r)
                             }
-                        } else if let Rvalue::Discriminant(d) = rvalue {
-                            trace!("skipping store of discriminant value {:?} to {:?}", d, dest);
+                            Rvalue::Repeat(op, c) => can_skip_operand(op) && can_skip_constant(c),
+                            Rvalue::AddressOf(_, _) => true,
+                            Rvalue::Len(_) => true,
+                            Rvalue::UnaryOp(_, op) => can_skip_operand(op),
+                            Rvalue::Aggregate(_, operands) => operands.iter().all(can_skip_operand),
+
+                            _ => false,
+                        };
+
+                        if can_skip {
+                            trace!("skipping store of {:?} to {:?}", rvalue, dest);
                             return;
                         }
                     }
                 }
             }
         }
 
-        self.locals.insert(*local);
+        self.local_counts[*local] += 1;
     }
 }
 
-struct LocalUpdater<'tcx> {
-    map: IndexVec<Local, Option<Local>>,
+struct StatementDeclMarker<'a, 'tcx> {
+    used_locals: IndexVec<Local, usize>,
+    statement: &'a Statement<'tcx>,
+}
+
+impl<'a, 'tcx> StatementDeclMarker<'a, 'tcx> {
+    pub fn new(local_count: usize, statement: &'a Statement<'tcx>) -> Self {
+        Self { used_locals: IndexVec::from_elem_n(0, local_count), statement }
+    }
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for StatementDeclMarker<'a, 'tcx> {
+    fn visit_local(&mut self, local: &Local, context: PlaceContext, _location: Location) {
+        // Skip the lvalue for assignments
+        if let StatementKind::Assign(box (p, _)) = self.statement.kind {
+            if p.local == *local && context.is_place_assignment() {
+                return;
+            }
+        }
+
+        self.used_locals[*local] += 1;
+    }
+}
+
+struct RemoveStatements<'a, 'tcx> {
+    used_locals: &'a mut IndexVec<Local, usize>,
+    arg_count: usize,
     tcx: TyCtxt<'tcx>,
+    modified: bool,
 }
 
-impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
+impl<'a, 'tcx> RemoveStatements<'a, 'tcx> {
+    fn new(
+        used_locals: &'a mut IndexVec<Local, usize>,
+        arg_count: usize,
+        tcx: TyCtxt<'tcx>,
+    ) -> Self {
+        Self { used_locals, arg_count, tcx, modified: false }
+    }
+
+    fn keep_local(&self, l: Local) -> bool {
+        trace!("keep_local({:?}): count: {:?}", l, self.used_locals[l]);
+        l.as_usize() <= self.arg_count || self.used_locals[l] != 0
+    }
+}
+
+impl<'a, 'tcx> MutVisitor<'tcx> for RemoveStatements<'a, 'tcx> {
     fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }
 
     fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
         // Remove unnecessary StorageLive and StorageDead annotations.
-        data.statements.retain(|stmt| match &stmt.kind {
-            StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => self.map[*l].is_some(),
-            StatementKind::Assign(box (place, _)) => self.map[place.local].is_some(),
-            _ => true,
+        let mut i = 0usize;
+        data.statements.retain(|stmt| {
+            let keep = match &stmt.kind {
+                StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
+                    self.keep_local(*l)
+                }
+                StatementKind::Assign(box (place, _)) => self.keep_local(place.local),
+                _ => true,
+            };
+
+            if !keep {
+                trace!("removing statement {:?}", stmt);
+                self.modified = true;
+
+                let mut visitor = StatementDeclMarker::new(self.used_locals.len(), stmt);
+                visitor.visit_statement(stmt, Location { block, statement_index: i });
+
+                for (local, count) in visitor.used_locals.iter_enumerated() {
+                    let used_count = &mut self.used_locals[local];
+
+                    // If this is the local we're removing...
+                    if *used_count != 0 {
+                        *used_count -= count;
+                    }
+                }
+            }
+
+            i += 1;
+
+            keep
         });
+
         self.super_basic_block_data(block, data);
     }
+}
+
+struct LocalUpdater<'tcx> {
+    map: IndexVec<Local, Option<Local>>,
+    tcx: TyCtxt<'tcx>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
 
     fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
         *l = self.map[*l].unwrap();

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

@@ -30,30 +30,30 @@ fn main() -> () {
 }
 
 alloc0 (static: FOO, size: 4, align: 4) {
-    ╾alloc10+0╼                                     │ ╾──╼
+    ╾alloc9+0─╼                                     │ ╾──╼
 }
 
-alloc10 (size: 168, align: 1) {
+alloc9 (size: 168, align: 1) {
     0x00 │ ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab │ ................
-    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾alloc5+0─╼ │ ............╾──╼
+    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾alloc4+0─╼ │ ............╾──╼
     0x20 │ 01 ef cd ab 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x30 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x40 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x50 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x60 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x70 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
-    0x80 │ 00 00 00 00 00 00 00 00 00 00 ╾alloc7+0─╼ 00 00 │ ..........╾──╼..
-    0x90 │ ╾alloc8+99╼ 00 00 00 00 00 00 00 00 00 00 00 00 │ ╾──╼............
+    0x80 │ 00 00 00 00 00 00 00 00 00 00 ╾alloc6+0─╼ 00 00 │ ..........╾──╼..
+    0x90 │ ╾alloc7+99╼ 00 00 00 00 00 00 00 00 00 00 00 00 │ ╾──╼............
     0xa0 │ 00 00 00 00 00 00 00 00                         │ ........
 }
 
-alloc5 (size: 4, align: 4) {
+alloc4 (size: 4, align: 4) {
     2a 00 00 00                                     │ *...
 }
 
-alloc7 (fn: main)
+alloc6 (fn: main)
 
-alloc8 (size: 100, align: 1) {
+alloc7 (size: 100, align: 1) {
     0x00 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x10 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     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

@@ -30,31 +30,31 @@ fn main() -> () {
 }
 
 alloc0 (static: FOO, size: 8, align: 8) {
-    ╾──────alloc10+0──────╼                         │ ╾──────╼
+    ╾──────alloc9+0───────╼                         │ ╾──────╼
 }
 
-alloc10 (size: 180, align: 1) {
+alloc9 (size: 180, align: 1) {
     0x00 │ ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab ab │ ................
-    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾─alloc5+0─ │ ............╾───
+    0x10 │ ab ab ab ab ab ab ab ab ab ab ab ab ╾─alloc4+0─ │ ............╾───
     0x20 │ ──────────╼ 01 ef cd ab 00 00 00 00 00 00 00 00 │ ───╼............
     0x30 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x40 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x50 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x60 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x70 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x80 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 ╾──── │ ..............╾─
-    0x90 │ ────alloc7+0────╼ 00 00 ╾──────alloc8+99──────╼ │ ─────╼..╾──────╼
+    0x90 │ ────alloc6+0────╼ 00 00 ╾──────alloc7+99──────╼ │ ─────╼..╾──────╼
     0xa0 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0xb0 │ 00 00 00 00                                     │ ....
 }
 
-alloc5 (size: 4, align: 4) {
+alloc4 (size: 4, align: 4) {
     2a 00 00 00                                     │ *...
 }
 
-alloc7 (fn: main)
+alloc6 (fn: main)
 
-alloc8 (size: 100, align: 1) {
+alloc7 (size: 100, align: 1) {
     0x00 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x10 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................
     0x20 │ 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 │ ................