Add RequiresStorage pass to decide which locals to save in generators

This avoids reserving storage in generators for locals that are moved
out of (and not re-initialized) prior to yield points.

4 files changed, 165 insertions, 26 deletions

diff --git a/src/librustc_mir/dataflow/at_location.rs b/src/librustc_mir/dataflow/at_location.rs
index f0014602e2d..7f5c77007ae 100644
--- a/src/librustc_mir/dataflow/at_location.rs
+++ b/src/librustc_mir/dataflow/at_location.rs
@@ -61,6 +61,7 @@ pub trait FlowsAtLocation {
 /// (e.g., via `reconstruct_statement_effect` and
 /// `reconstruct_terminator_effect`; don't forget to call
 /// `apply_local_effect`).
+#[derive(Clone)]
 pub struct FlowAtLocation<'tcx, BD, DR = DataflowResults<'tcx, BD>>
 where
     BD: BitDenotation<'tcx>,
diff --git a/src/librustc_mir/dataflow/impls/storage_liveness.rs b/src/librustc_mir/dataflow/impls/storage_liveness.rs
index d2003993d45..64d0ae016a3 100644
--- a/src/librustc_mir/dataflow/impls/storage_liveness.rs
+++ b/src/librustc_mir/dataflow/impls/storage_liveness.rs
@@ -1,7 +1,13 @@
 pub use super::*;
 
 use rustc::mir::*;
+use rustc::mir::visit::{
+    PlaceContext, Visitor, NonMutatingUseContext,
+};
+use std::cell::RefCell;
 use crate::dataflow::BitDenotation;
+use crate::dataflow::HaveBeenBorrowedLocals;
+use crate::dataflow::{DataflowResults, DataflowResultsCursor, DataflowResultsRefCursor};
 
 #[derive(Copy, Clone)]
 pub struct MaybeStorageLive<'a, 'tcx> {
@@ -63,3 +69,126 @@ impl<'a, 'tcx> BottomValue for MaybeStorageLive<'a, 'tcx> {
     /// bottom = dead
     const BOTTOM_VALUE: bool = false;
 }
+
+/// Dataflow analysis that determines whether each local requires storage at a
+/// given location; i.e. whether its storage can go away without being observed.
+///
+/// In the case of a movable generator, borrowed_locals can be `None` and we
+/// will not consider borrows in this pass. This relies on the fact that we only
+/// use this pass at yield points for these generators.
+#[derive(Clone)]
+pub struct RequiresStorage<'mir, 'tcx, 'b> {
+    body: &'mir Body<'tcx>,
+    borrowed_locals:
+        RefCell<DataflowResultsRefCursor<'mir, 'tcx, 'b, HaveBeenBorrowedLocals<'mir, 'tcx>>>,
+}
+
+impl<'mir, 'tcx: 'mir, 'b> RequiresStorage<'mir, 'tcx, 'b> {
+    pub fn new(
+        body: &'mir Body<'tcx>,
+        borrowed_locals: &'b DataflowResults<'tcx, HaveBeenBorrowedLocals<'mir, 'tcx>>,
+    ) -> Self {
+        RequiresStorage {
+            body,
+            borrowed_locals: RefCell::new(DataflowResultsCursor::new(borrowed_locals, body)),
+        }
+    }
+
+    pub fn body(&self) -> &Body<'tcx> {
+        self.body
+    }
+}
+
+impl<'mir, 'tcx, 'b> BitDenotation<'tcx> for RequiresStorage<'mir, 'tcx, 'b> {
+    type Idx = Local;
+    fn name() -> &'static str { "requires_storage" }
+    fn bits_per_block(&self) -> usize {
+        self.body.local_decls.len()
+    }
+
+    fn start_block_effect(&self, _sets: &mut BitSet<Local>) {
+        // Nothing is live on function entry
+    }
+
+    fn statement_effect(&self,
+                        sets: &mut GenKillSet<Local>,
+                        loc: Location) {
+        self.check_for_move(sets, loc);
+        self.check_for_borrow(sets, loc);
+
+        let stmt = &self.body[loc.block].statements[loc.statement_index];
+        match stmt.kind {
+            StatementKind::StorageLive(l) => sets.gen(l),
+            StatementKind::StorageDead(l) => sets.kill(l),
+            StatementKind::Assign(ref place, _)
+            | StatementKind::SetDiscriminant { ref place, .. } => {
+                place.base_local().map(|l| sets.gen(l));
+            }
+            StatementKind::InlineAsm(box InlineAsm { ref outputs, .. }) => {
+                for p in &**outputs {
+                    p.base_local().map(|l| sets.gen(l));
+                }
+            }
+            _ => (),
+        }
+    }
+
+    fn terminator_effect(&self,
+                         sets: &mut GenKillSet<Local>,
+                         loc: Location) {
+        self.check_for_move(sets, loc);
+        self.check_for_borrow(sets, loc);
+    }
+
+    fn propagate_call_return(
+        &self,
+        in_out: &mut BitSet<Local>,
+        _call_bb: mir::BasicBlock,
+        _dest_bb: mir::BasicBlock,
+        dest_place: &mir::Place<'tcx>,
+    ) {
+        dest_place.base_local().map(|l| in_out.insert(l));
+    }
+}
+
+impl<'mir, 'tcx, 'b> RequiresStorage<'mir, 'tcx, 'b> {
+    /// Kill locals that are fully moved and have not been borrowed.
+    fn check_for_move(&self, sets: &mut GenKillSet<Local>, loc: Location) {
+        let mut visitor = MoveVisitor {
+            sets,
+            borrowed_locals: &self.borrowed_locals,
+        };
+        visitor.visit_location(self.body, loc);
+    }
+
+    /// Gen locals that are newly borrowed. This includes borrowing any part of
+    /// a local (we rely on this behavior of `HaveBeenBorrowedLocals`).
+    fn check_for_borrow(&self, sets: &mut GenKillSet<Local>, loc: Location) {
+        let mut borrowed_locals = self.borrowed_locals.borrow_mut();
+        borrowed_locals.seek(loc);
+        borrowed_locals.each_gen_bit(|l| sets.gen(l));
+    }
+}
+
+impl<'mir, 'tcx, 'b> BottomValue for RequiresStorage<'mir, 'tcx, 'b> {
+    /// bottom = dead
+    const BOTTOM_VALUE: bool = false;
+}
+
+struct MoveVisitor<'a, 'b, 'mir, 'tcx> {
+    borrowed_locals:
+        &'a RefCell<DataflowResultsRefCursor<'mir, 'tcx, 'b, HaveBeenBorrowedLocals<'mir, 'tcx>>>,
+    sets: &'a mut GenKillSet<Local>,
+}
+
+impl<'a, 'b, 'mir: 'a, 'tcx> Visitor<'tcx> for MoveVisitor<'a, 'b, 'mir, 'tcx> {
+    fn visit_local(&mut self, local: &Local, context: PlaceContext, loc: Location) {
+        if PlaceContext::NonMutatingUse(NonMutatingUseContext::Move) == context {
+            let mut borrowed_locals = self.borrowed_locals.borrow_mut();
+            borrowed_locals.seek(loc);
+            if !borrowed_locals.contains(*local) {
+                self.sets.kill(*local);
+            }
+        }
+    }
+}
diff --git a/src/librustc_mir/dataflow/mod.rs b/src/librustc_mir/dataflow/mod.rs
index fccd1fdf8f2..060b2be5cd9 100644
--- a/src/librustc_mir/dataflow/mod.rs
+++ b/src/librustc_mir/dataflow/mod.rs
@@ -17,7 +17,7 @@ use std::io;
 use std::path::PathBuf;
 use std::usize;
 
-pub use self::impls::{MaybeStorageLive};
+pub use self::impls::{MaybeStorageLive, RequiresStorage};
 pub use self::impls::{MaybeInitializedPlaces, MaybeUninitializedPlaces};
 pub use self::impls::DefinitelyInitializedPlaces;
 pub use self::impls::EverInitializedPlaces;
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 3d0d565c53b..54f69adb69c 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -66,9 +66,9 @@ use std::mem;
 use crate::transform::{MirPass, MirSource};
 use crate::transform::simplify;
 use crate::transform::no_landing_pads::no_landing_pads;
-use crate::dataflow::{DataflowResults, DataflowResultsConsumer, FlowAtLocation, FlowAtLocationOwned};
+use crate::dataflow::{DataflowResults, DataflowResultsConsumer, FlowAtLocation};
 use crate::dataflow::{do_dataflow, DebugFormatted, state_for_location};
-use crate::dataflow::{MaybeStorageLive, HaveBeenBorrowedLocals};
+use crate::dataflow::{MaybeStorageLive, HaveBeenBorrowedLocals, RequiresStorage};
 use crate::util::dump_mir;
 use crate::util::liveness;
 
@@ -437,16 +437,17 @@ fn locals_live_across_suspend_points(
 
     // Calculate the MIR locals which have been previously
     // borrowed (even if they are still active).
-    // This is only used for immovable generators.
-    let borrowed_locals = if !movable {
-        let analysis = HaveBeenBorrowedLocals::new(body);
-        let result =
-            do_dataflow(tcx, body, def_id, &[], &dead_unwinds, analysis,
-                        |bd, p| DebugFormatted::new(&bd.body().local_decls[p]));
-        Some((analysis, result))
-    } else {
-        None
-    };
+    let borrowed_locals_analysis = HaveBeenBorrowedLocals::new(body);
+    let borrowed_locals_result =
+        do_dataflow(tcx, body, def_id, &[], &dead_unwinds, borrowed_locals_analysis,
+                    |bd, p| DebugFormatted::new(&bd.body().local_decls[p]));
+
+    // Calculate the MIR locals that we actually need to keep storage around
+    // for.
+    let requires_storage_analysis = RequiresStorage::new(body, &borrowed_locals_result);
+    let requires_storage =
+        do_dataflow(tcx, body, def_id, &[], &dead_unwinds, requires_storage_analysis.clone(),
+                    |bd, p| DebugFormatted::new(&bd.body().local_decls[p]));
 
     // Calculate the liveness of MIR locals ignoring borrows.
     let mut live_locals = liveness::LiveVarSet::new_empty(body.local_decls.len());
@@ -471,10 +472,10 @@ fn locals_live_across_suspend_points(
                 statement_index: data.statements.len(),
             };
 
-            if let Some((ref analysis, ref result)) = borrowed_locals {
+            if !movable {
                 let borrowed_locals = state_for_location(loc,
-                                                         analysis,
-                                                         result,
+                                                         &borrowed_locals_analysis,
+                                                         &borrowed_locals_result,
                                                          body);
                 // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
                 // This is correct for movable generators since borrows cannot live across
@@ -489,27 +490,34 @@ fn locals_live_across_suspend_points(
                 liveness.outs[block].union(&borrowed_locals);
             }
 
-            let mut storage_liveness = state_for_location(loc,
-                                                          &storage_live_analysis,
-                                                          &storage_live,
-                                                          body);
+            let storage_liveness = state_for_location(loc,
+                                                      &storage_live_analysis,
+                                                      &storage_live,
+                                                      body);
 
             // Store the storage liveness for later use so we can restore the state
             // after a suspension point
             storage_liveness_map.insert(block, storage_liveness.clone());
 
-            // Mark locals without storage statements as always having live storage
-            storage_liveness.union(&ignored.0);
+            let mut storage_required = state_for_location(loc,
+                                                          &requires_storage_analysis,
+                                                          &requires_storage,
+                                                          body);
+
+            // Mark locals without storage statements as always requiring storage
+            storage_required.union(&ignored.0);
 
             // Locals live are live at this point only if they are used across
             // suspension points (the `liveness` variable)
-            // and their storage is live (the `storage_liveness` variable)
-            let mut live_locals_here = storage_liveness;
+            // and their storage is required (the `storage_required` variable)
+            let mut live_locals_here = storage_required;
             live_locals_here.intersect(&liveness.outs[block]);
 
             // The generator argument is ignored
             live_locals_here.remove(self_arg());
 
+            debug!("loc = {:?}, live_locals_here = {:?}", loc, live_locals_here);
+
             // Add the locals live at this suspension point to the set of locals which live across
             // any suspension points
             live_locals.union(&live_locals_here);
@@ -517,6 +525,7 @@ fn locals_live_across_suspend_points(
             live_locals_at_suspension_points.push(live_locals_here);
         }
     }
+    debug!("live_locals = {:?}", live_locals);
 
     // Renumber our liveness_map bitsets to include only the locals we are
     // saving.
@@ -627,7 +636,7 @@ struct StorageConflictVisitor<'body, 'tcx, 's> {
 impl<'body, 'tcx, 's> DataflowResultsConsumer<'body, 'tcx>
     for StorageConflictVisitor<'body, 'tcx, 's>
 {
-    type FlowState = FlowAtLocationOwned<'tcx, MaybeStorageLive<'body, 'tcx>>;
+    type FlowState = FlowAtLocation<'tcx, MaybeStorageLive<'body, 'tcx>>;
 
     fn body(&self) -> &'body Body<'tcx> {
         self.body
@@ -657,7 +666,7 @@ impl<'body, 'tcx, 's> DataflowResultsConsumer<'body, 'tcx>
 
 impl<'body, 'tcx, 's> StorageConflictVisitor<'body, 'tcx, 's> {
     fn apply_state(&mut self,
-                   flow_state: &FlowAtLocationOwned<'tcx, MaybeStorageLive<'body, 'tcx>>,
+                   flow_state: &FlowAtLocation<'tcx, MaybeStorageLive<'body, 'tcx>>,
                    loc: Location) {
         // Ignore unreachable blocks.
         match self.body.basic_blocks()[loc.block].terminator().kind {