A way to remove otherwise unused locals from MIR

Replaces the hack where a similar thing is done within trans.

6 files changed, 137 insertions, 35 deletions

diff --git a/src/librustc/mir/visit.rs b/src/librustc/mir/visit.rs
index db7267ca0d4..b5da304a109 100644
--- a/src/librustc/mir/visit.rs
+++ b/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,
diff --git a/src/librustc_data_structures/indexed_vec.rs b/src/librustc_data_structures/indexed_vec.rs
index 91234631499..00cea9cbdf6 100644
--- a/src/librustc_data_structures/indexed_vec.rs
+++ b/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> {
diff --git a/src/librustc_driver/driver.rs b/src/librustc_driver/driver.rs
index 318616726f4..ea1728daf80 100644
--- a/src/librustc_driver/driver.rs
+++ b/src/librustc_driver/driver.rs
@@ -915,17 +915,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);
         });
@@ -987,13 +989,13 @@ 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);
@@ -1001,13 +1003,14 @@ pub fn phase_4_translate_to_llvm<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         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"));
 
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
index 7bcb89b5895..ae255f70fb7 100644
--- a/src/librustc_mir/transform/mod.rs
+++ b/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;
diff --git a/src/librustc_mir/transform/simplify_cfg.rs b/src/librustc_mir/transform/simplify.rs
index 1a8a5fa18cf..d5fc90289e2 100644
--- a/src/librustc_mir/transform/simplify_cfg.rs
+++ b/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);
+    }
+}
diff --git a/src/librustc_trans/mir/analyze.rs b/src/librustc_trans/mir/analyze.rs
index a934da12b9e..e4d0533ec87 100644
--- a/src/librustc_trans/mir/analyze.rs
+++ b/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