Auto merge of #36306 - nagisa:mir-local-cleanup, r=eddyb

A way to remove otherwise unused locals from MIR

There is a certain amount of desire for a pass which cleans up the provably unused variables (no assignments or reads). There has been an implementation of such pass by @scottcarr, and another (two!) implementations by me in my own dataflow efforts.

PR like #35916 proves that this pass is useful even on its own, which is why I cherry-picked it out from my dataflow effort.

@nikomatsakis previously expressed concerns over this pass not seeming to be very cheap to run and therefore unsuitable for regular cleanup duties. Turns out, regular cleanup of local declarations is not at all necessary, at least now, because majority of passes simply do not (or should not) care about them. That’s why it is viable to only run this pass once (perhaps a few more times in the future?) per function, right before translation.

r? @eddyb or @nikomatsakis

src/librustc/mir/visit.rs

@@ -157,14 +157,14 @@ macro_rules! make_mir_visitor {
 
             fn visit_projection(&mut self,
                                 lvalue: & $($mutability)* LvalueProjection<'tcx>,
-                                context: LvalueContext,
+                                context: LvalueContext<'tcx>,
                                 location: Location) {
                 self.super_projection(lvalue, context, location);
             }
 
             fn visit_projection_elem(&mut self,
                                      lvalue: & $($mutability)* LvalueElem<'tcx>,
-                                     context: LvalueContext,
+                                     context: LvalueContext<'tcx>,
                                      location: Location) {
                 self.super_projection_elem(lvalue, context, location);
             }
@@ -579,7 +579,7 @@ macro_rules! make_mir_visitor {
 
             fn super_projection(&mut self,
                                 proj: & $($mutability)* LvalueProjection<'tcx>,
-                                context: LvalueContext,
+                                context: LvalueContext<'tcx>,
                                 location: Location) {
                 let Projection {
                     ref $($mutability)* base,
@@ -596,7 +596,7 @@ macro_rules! make_mir_visitor {
 
             fn super_projection_elem(&mut self,
                                      proj: & $($mutability)* LvalueElem<'tcx>,
-                                     _context: LvalueContext,
+                                     _context: LvalueContext<'tcx>,
                                      location: Location) {
                 match *proj {
                     ProjectionElem::Deref => {
@@ -739,7 +739,7 @@ macro_rules! make_mir_visitor {
 make_mir_visitor!(Visitor,);
 make_mir_visitor!(MutVisitor,mut);
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum LvalueContext<'tcx> {
     // Appears as LHS of an assignment
     Store,

src/librustc_data_structures/indexed_vec.rs

@@ -149,6 +149,21 @@ impl<I: Idx, T> IndexVec<I, T> {
     pub fn last(&self) -> Option<I> {
         self.len().checked_sub(1).map(I::new)
     }
+
+    #[inline]
+    pub fn shrink_to_fit(&mut self) {
+        self.raw.shrink_to_fit()
+    }
+
+    #[inline]
+    pub fn swap(&mut self, a: usize, b: usize) {
+        self.raw.swap(a, b)
+    }
+
+    #[inline]
+    pub fn truncate(&mut self, a: usize) {
+        self.raw.truncate(a)
+    }
 }
 
 impl<I: Idx, T> Index<I> for IndexVec<I, T> {

src/librustc_driver/driver.rs

@@ -917,17 +917,19 @@ pub fn phase_3_run_analysis_passes<'tcx, F, R>(sess: &'tcx Session,
              "MIR dump",
              || mir::mir_map::build_mir_for_crate(tcx));
 
-        time(time_passes, "MIR passes", || {
+        time(time_passes, "MIR cleanup and validation", || {
             let mut passes = sess.mir_passes.borrow_mut();
-            // Push all the built-in passes.
+            // Push all the built-in validation passes.
+            // NB: if you’re adding an *optimisation* it ought to go to another set of passes
+            // in stage 4 below.
             passes.push_hook(box mir::transform::dump_mir::DumpMir);
-            passes.push_pass(box mir::transform::simplify_cfg::SimplifyCfg::new("initial"));
+            passes.push_pass(box mir::transform::simplify::SimplifyCfg::new("initial"));
             passes.push_pass(
                 box mir::transform::qualify_consts::QualifyAndPromoteConstants::default());
             passes.push_pass(box mir::transform::type_check::TypeckMir);
             passes.push_pass(
                 box mir::transform::simplify_branches::SimplifyBranches::new("initial"));
-            passes.push_pass(box mir::transform::simplify_cfg::SimplifyCfg::new("qualify-consts"));
+            passes.push_pass(box mir::transform::simplify::SimplifyCfg::new("qualify-consts"));
             // And run everything.
             passes.run_passes(tcx);
         });
@@ -989,27 +991,28 @@ pub fn phase_4_translate_to_llvm<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
          "resolving dependency formats",
          || dependency_format::calculate(&tcx.sess));
 
-    // Run the passes that transform the MIR into a more suitable for translation
-    // to LLVM code.
-    time(time_passes, "Prepare MIR codegen passes", || {
+    // Run the passes that transform the MIR into a more suitable form for translation to LLVM
+    // code.
+    time(time_passes, "MIR optimisations", || {
         let mut passes = ::rustc::mir::transform::Passes::new();
         passes.push_hook(box mir::transform::dump_mir::DumpMir);
         passes.push_pass(box mir::transform::no_landing_pads::NoLandingPads);
-        passes.push_pass(box mir::transform::simplify_cfg::SimplifyCfg::new("no-landing-pads"));
+        passes.push_pass(box mir::transform::simplify::SimplifyCfg::new("no-landing-pads"));
 
         // From here on out, regions are gone.
         passes.push_pass(box mir::transform::erase_regions::EraseRegions);
 
         passes.push_pass(box mir::transform::add_call_guards::AddCallGuards);
         passes.push_pass(box borrowck::ElaborateDrops);
         passes.push_pass(box mir::transform::no_landing_pads::NoLandingPads);
-        passes.push_pass(box mir::transform::simplify_cfg::SimplifyCfg::new("elaborate-drops"));
+        passes.push_pass(box mir::transform::simplify::SimplifyCfg::new("elaborate-drops"));
 
         // No lifetime analysis based on borrowing can be done from here on out.
         passes.push_pass(box mir::transform::instcombine::InstCombine::new());
         passes.push_pass(box mir::transform::deaggregator::Deaggregator);
         passes.push_pass(box mir::transform::copy_prop::CopyPropagation);
 
+        passes.push_pass(box mir::transform::simplify::SimplifyLocals);
         passes.push_pass(box mir::transform::add_call_guards::AddCallGuards);
         passes.push_pass(box mir::transform::dump_mir::Marker("PreTrans"));
 

src/librustc_mir/transform/mod.rs

@@ -9,7 +9,7 @@
 // except according to those terms.
 
 pub mod simplify_branches;
-pub mod simplify_cfg;
+pub mod simplify;
 pub mod erase_regions;
 pub mod no_landing_pads;
 pub mod type_check;

src/librustc_mir/transform/simplify_cfg.rs → src/librustc_mir/transform/simplify.rs

@@ -8,35 +8,41 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-//! A pass that removes various redundancies in the CFG. It should be
-//! called after every significant CFG modification to tidy things
-//! up.
+//! A number of passes which remove various redundancies in the CFG.
 //!
-//! This pass must also be run before any analysis passes because it removes
-//! dead blocks, and some of these can be ill-typed.
+//! The `SimplifyCfg` pass gets rid of unnecessary blocks in the CFG, whereas the `SimplifyLocals`
+//! gets rid of all the unnecessary local variable declarations.
 //!
-//! The cause of that is that typeck lets most blocks whose end is not
-//! reachable have an arbitrary return type, rather than having the
-//! usual () return type (as a note, typeck's notion of reachability
-//! is in fact slightly weaker than MIR CFG reachability - see #31617).
+//! The `SimplifyLocals` pass is kinda expensive and therefore not very suitable to be run often.
+//! Most of the passes should not care or be impacted in meaningful ways due to extra locals
+//! either, so running the pass once, right before translation, should suffice.
+//!
+//! On the other side of the spectrum, the `SimplifyCfg` pass is considerably cheap to run, thus
+//! one should run it after every pass which may modify CFG in significant ways. This pass must
+//! also be run before any analysis passes because it removes dead blocks, and some of these can be
+//! ill-typed.
+//!
+//! The cause of this typing issue is typeck allowing most blocks whose end is not reachable have
+//! an arbitrary return type, rather than having the usual () return type (as a note, typeck's
+//! notion of reachability is in fact slightly weaker than MIR CFG reachability - see #31617). A
+//! standard example of the situation is:
 //!
-//! A standard example of the situation is:
 //! ```rust
 //!   fn example() {
 //!       let _a: char = { return; };
 //!   }
 //! ```
 //!
-//! Here the block (`{ return; }`) has the return type `char`,
-//! rather than `()`, but the MIR we naively generate still contains
-//! the `_a = ()` write in the unreachable block "after" the return.
-
+//! Here the block (`{ return; }`) has the return type `char`, rather than `()`, but the MIR we
+//! naively generate still contains the `_a = ()` write in the unreachable block "after" the
+//! return.
 
 use rustc_data_structures::bitvec::BitVector;
 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
 use rustc::ty::TyCtxt;
 use rustc::mir::*;
 use rustc::mir::transform::{MirPass, MirSource, Pass};
+use rustc::mir::visit::{MutVisitor, Visitor, LvalueContext};
 use std::fmt;
 
 pub struct SimplifyCfg<'a> { label: &'a str }
@@ -257,3 +263,87 @@ fn remove_dead_blocks(mir: &mut Mir) {
         }
     }
 }
+
+
+pub struct SimplifyLocals;
+
+impl Pass for SimplifyLocals {
+    fn name(&self) -> ::std::borrow::Cow<'static, str> { "SimplifyLocals".into() }
+}
+
+impl<'tcx> MirPass<'tcx> for SimplifyLocals {
+    fn run_pass<'a>(&mut self, _: TyCtxt<'a, 'tcx, 'tcx>, _: MirSource, mir: &mut Mir<'tcx>) {
+        let mut marker = DeclMarker { locals: BitVector::new(mir.local_decls.len()) };
+        marker.visit_mir(mir);
+        // Return pointer and arguments are always live
+        marker.locals.insert(0);
+        for idx in mir.args_iter() {
+            marker.locals.insert(idx.index());
+        }
+        let map = make_local_map(&mut mir.local_decls, marker.locals);
+        // Update references to all vars and tmps now
+        LocalUpdater { map: map }.visit_mir(mir);
+        mir.local_decls.shrink_to_fit();
+    }
+}
+
+/// Construct the mapping while swapping out unused stuff out from the `vec`.
+fn make_local_map<'tcx, I: Idx, V>(vec: &mut IndexVec<I, V>, mask: BitVector) -> Vec<usize> {
+    let mut map: Vec<usize> = ::std::iter::repeat(!0).take(vec.len()).collect();
+    let mut used = 0;
+    for alive_index in mask.iter() {
+        map[alive_index] = used;
+        if alive_index != used {
+            vec.swap(alive_index, used);
+        }
+        used += 1;
+    }
+    vec.truncate(used);
+    map
+}
+
+struct DeclMarker {
+    pub locals: BitVector,
+}
+
+impl<'tcx> Visitor<'tcx> for DeclMarker {
+    fn visit_lvalue(&mut self, lval: &Lvalue<'tcx>, ctx: LvalueContext<'tcx>, loc: Location) {
+        if ctx == LvalueContext::StorageLive || ctx == LvalueContext::StorageDead {
+            // ignore these altogether, they get removed along with their otherwise unused decls.
+            return;
+        }
+        if let Lvalue::Local(ref v) = *lval {
+            self.locals.insert(v.index());
+        }
+        self.super_lvalue(lval, ctx, loc);
+    }
+}
+
+struct LocalUpdater {
+    map: Vec<usize>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for LocalUpdater {
+    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(ref lval) | StatementKind::StorageDead(ref lval) => {
+                    match *lval {
+                        Lvalue::Local(l) => self.map[l.index()] != !0,
+                        _ => true
+                    }
+                }
+                _ => true
+            }
+        });
+        self.super_basic_block_data(block, data);
+    }
+    fn visit_lvalue(&mut self, lval: &mut Lvalue<'tcx>, ctx: LvalueContext<'tcx>, loc: Location) {
+        match *lval {
+            Lvalue::Local(ref mut l) => *l = Local::new(self.map[l.index()]),
+            _ => (),
+        };
+        self.super_lvalue(lval, ctx, loc);
+    }
+}

src/librustc_trans/mir/analyze.rs

@@ -42,12 +42,6 @@ pub fn lvalue_locals<'bcx, 'tcx>(bcx: Block<'bcx,'tcx>,
                     common::type_is_fat_ptr(bcx.tcx(), ty));
         } else if common::type_is_imm_pair(bcx.ccx(), ty) {
             // We allow pairs and uses of any of their 2 fields.
-        } else if !analyzer.seen_assigned.contains(index) {
-            // No assignment has been seen, which means that
-            // either the local has been marked as lvalue
-            // already, or there is no possible initialization
-            // for the local, making any reads invalid.
-            // This is useful in weeding out dead temps.
         } else {
             // These sorts of types require an alloca. Note that
             // type_is_immediate() may *still* be true, particularly