Force locals to be live after they are borrowed for immovable generators. Fixes #47736

src/librustc_mir/dataflow/impls/borrowed_locals.rs

@@ -0,0 +1,118 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+pub use super::*;
+
+use rustc::mir::*;
+use rustc::mir::visit::Visitor;
+use dataflow::BitDenotation;
+
+/// This calculates if any part of a MIR local could have previously been borrowed.
+/// This means that once a local has been borrowed, its bit will always be set
+/// from that point and onwards, even if the borrow ends. You could also think of this
+/// as computing the lifetimes of infinite borrows.
+/// This is used to compute which locals are live during a yield expression for
+/// immovable generators.
+#[derive(Copy, Clone)]
+pub struct HaveBeenBorrowedLocals<'a, 'tcx: 'a> {
+    mir: &'a Mir<'tcx>,
+}
+
+impl<'a, 'tcx: 'a> HaveBeenBorrowedLocals<'a, 'tcx> {
+    pub fn new(mir: &'a Mir<'tcx>)
+               -> Self {
+        HaveBeenBorrowedLocals { mir: mir }
+    }
+
+    pub fn mir(&self) -> &Mir<'tcx> {
+        self.mir
+    }
+}
+
+impl<'a, 'tcx> BitDenotation for HaveBeenBorrowedLocals<'a, 'tcx> {
+    type Idx = Local;
+    fn name() -> &'static str { "has_been_borrowed_locals" }
+    fn bits_per_block(&self) -> usize {
+        self.mir.local_decls.len()
+    }
+
+    fn start_block_effect(&self, _sets: &mut IdxSet<Local>) {
+        // Nothing is borrowed on function entry
+    }
+
+    fn statement_effect(&self,
+                        sets: &mut BlockSets<Local>,
+                        loc: Location) {
+        BorrowedLocalsVisitor {
+            sets,
+        }.visit_statement(loc.block, &self.mir[loc.block].statements[loc.statement_index], loc);
+    }
+
+    fn terminator_effect(&self,
+                         sets: &mut BlockSets<Local>,
+                         loc: Location) {
+        BorrowedLocalsVisitor {
+            sets,
+        }.visit_terminator(loc.block, self.mir[loc.block].terminator(), loc);
+    }
+
+    fn propagate_call_return(&self,
+                             _in_out: &mut IdxSet<Local>,
+                             _call_bb: mir::BasicBlock,
+                             _dest_bb: mir::BasicBlock,
+                             _dest_place: &mir::Place) {
+        // Nothing to do when a call returns successfully
+    }
+}
+
+impl<'a, 'tcx> BitwiseOperator for HaveBeenBorrowedLocals<'a, 'tcx> {
+    #[inline]
+    fn join(&self, pred1: usize, pred2: usize) -> usize {
+        pred1 | pred2 // "maybe" means we union effects of both preds
+    }
+}
+
+impl<'a, 'tcx> InitialFlow for HaveBeenBorrowedLocals<'a, 'tcx> {
+    #[inline]
+    fn bottom_value() -> bool {
+        false // bottom = unborrowed
+    }
+}
+
+struct BorrowedLocalsVisitor<'b, 'c: 'b> {
+    sets: &'b mut BlockSets<'c, Local>,
+}
+
+fn find_local<'tcx>(place: &Place<'tcx>) -> Option<Local> {
+    match *place {
+        Place::Local(l) => Some(l),
+        Place::Static(..) => None,
+        Place::Projection(ref proj) => {
+            match proj.elem {
+                ProjectionElem::Deref => None,
+                _ => find_local(&proj.base)
+            }
+        }
+    }
+}
+
+impl<'tcx, 'b, 'c> Visitor<'tcx> for BorrowedLocalsVisitor<'b, 'c> {
+    fn visit_rvalue(&mut self,
+                    rvalue: &Rvalue<'tcx>,
+                    location: Location) {
+        if let Rvalue::Ref(_, _, ref place) = *rvalue {
+            if let Some(local) = find_local(place) {
+                self.sets.gen(&local);
+            }
+        }
+
+        self.super_rvalue(rvalue, location)
+    }
+}

src/librustc_mir/dataflow/impls/mod.rs

@@ -33,6 +33,10 @@ mod storage_liveness;
 
 pub use self::storage_liveness::*;
 
+mod borrowed_locals;
+
+pub use self::borrowed_locals::*;
+
 #[allow(dead_code)]
 pub(super) mod borrows;
 

src/librustc_mir/dataflow/mod.rs

@@ -30,6 +30,7 @@ pub use self::impls::{MaybeInitializedLvals, MaybeUninitializedLvals};
 pub use self::impls::{DefinitelyInitializedLvals, MovingOutStatements};
 pub use self::impls::EverInitializedLvals;
 pub use self::impls::borrows::{Borrows, BorrowData};
+pub use self::impls::HaveBeenBorrowedLocals;
 pub(crate) use self::impls::borrows::{ActiveBorrows, Reservations, ReserveOrActivateIndex};
 pub use self::at_location::{FlowAtLocation, FlowsAtLocation};
 pub(crate) use self::drop_flag_effects::*;

src/librustc_mir/transform/generator.rs

@@ -78,7 +78,8 @@ use std::mem;
 use transform::{MirPass, MirSource};
 use transform::simplify;
 use transform::no_landing_pads::no_landing_pads;
-use dataflow::{do_dataflow, DebugFormatted, MaybeStorageLive, state_for_location};
+use dataflow::{do_dataflow, DebugFormatted, state_for_location};
+use dataflow::{MaybeStorageLive, HaveBeenBorrowedLocals};
 
 pub struct StateTransform;
 
@@ -369,17 +370,33 @@ fn locals_live_across_suspend_points<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                                 HashMap<BasicBlock, liveness::LocalSet>) {
     let dead_unwinds = IdxSetBuf::new_empty(mir.basic_blocks().len());
     let node_id = tcx.hir.as_local_node_id(source.def_id).unwrap();
-    let analysis = MaybeStorageLive::new(mir);
+
+    // Calculate when MIR locals have live storage. This gives us an upper bound of their
+    // lifetimes.
+    let storage_live_analysis = MaybeStorageLive::new(mir);
     let storage_live =
-        do_dataflow(tcx, mir, node_id, &[], &dead_unwinds, analysis,
+        do_dataflow(tcx, mir, node_id, &[], &dead_unwinds, storage_live_analysis,
                     |bd, p| DebugFormatted::new(&bd.mir().local_decls[p]));
 
+    // Find the MIR locals which do not use StorageLive/StorageDead statements.
+    // The storage of these locals are always live.
     let mut ignored = StorageIgnored(IdxSetBuf::new_filled(mir.local_decls.len()));
     ignored.visit_mir(mir);
 
-    let mut borrowed_locals = BorrowedLocals(IdxSetBuf::new_empty(mir.local_decls.len()));
-    borrowed_locals.visit_mir(mir);
+    // 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(mir);
+        let result =
+            do_dataflow(tcx, mir, node_id, &[], &dead_unwinds, analysis,
+                        |bd, p| DebugFormatted::new(&bd.mir().local_decls[p]));
+        Some((analysis, result))
+    } else {
+        None
+    };
 
+    // Calculate the liveness of MIR locals ignoring borrows.
     let mut set = liveness::LocalSet::new_empty(mir.local_decls.len());
     let mut liveness = liveness::liveness_of_locals(mir, LivenessMode {
         include_regular_use: true,
@@ -396,24 +413,41 @@ fn locals_live_across_suspend_points<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                 statement_index: data.statements.len(),
             };
 
-            let storage_liveness = state_for_location(loc, &analysis, &storage_live, mir);
-
-            storage_liveness_map.insert(block, storage_liveness.clone());
-
-            let mut live_locals = storage_liveness;
-
-            // Mark locals without storage statements as always having live storage
-            live_locals.union(&ignored.0);
-
-            if !movable {
-                // For immovable generators we consider borrowed locals to always be live.
-                // This effectively makes those locals use just the storage liveness.
-                liveness.outs[block].union(&borrowed_locals.0);
+            if let Some((ref analysis, ref result)) = borrowed_locals {
+                let borrowed_locals = state_for_location(loc,
+                                                         analysis,
+                                                         result,
+                                                         mir);
+                // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
+                // This is correct for movable generators since borrows cannot live across
+                // suspension points. However for immovable generators we need to account for
+                // borrows, so we conseratively assume that all borrowed locals live forever.
+                // To do this we just union our `liveness` result with `borrowed_locals`, which
+                // contains all the locals which has been borrowed before this suspension point.
+                // If a borrow is converted to a raw reference, we must also assume that it lives
+                // forever. Note that the final liveness is still bounded by the storage liveness
+                // of the local, which happens using the `intersect` operation below.
+                liveness.outs[block].union(&borrowed_locals);
             }
 
-            // Locals live are live at this point only if they are used across suspension points
-            // and their storage is live
-            live_locals.intersect(&liveness.outs[block]);
+            let mut storage_liveness = state_for_location(loc,
+                                                          &storage_live_analysis,
+                                                          &storage_live,
+                                                          mir);
+
+            // 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);
+
+            // 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)
+            storage_liveness.intersect(&liveness.outs[block]);
+
+            let live_locals = storage_liveness;
 
             // Add the locals life at this suspension point to the set of locals which live across
             // any suspension points

src/test/run-pass/generator/too-live-local-in-immovable-gen.rs

@@ -0,0 +1,28 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(generators)]
+
+fn main() {
+    unsafe {
+        static move || {
+            // Tests that the generator transformation finds out that `a` is not live
+            // during the yield expression. Type checking will also compute liveness
+            // and it should also find out that `a` is not live.
+            // The compiler will panic if the generator transformation finds that
+            // `a` is live and type checking finds it dead.
+            let a = {
+                yield ();
+                4i32
+            };
+            &a;
+        };
+    }
+}