Generator bugfixes

src/librustc/middle/region.rs

@@ -467,9 +467,13 @@ impl<'tcx> Visitor<'tcx> for ExprLocatorVisitor {
     }
 
     fn visit_pat(&mut self, pat: &'tcx Pat) {
+        intravisit::walk_pat(self, pat);
+
         self.expr_and_pat_count += 1;
 
-        intravisit::walk_pat(self, pat);
+        if pat.id == self.id {
+            self.result = Some(self.expr_and_pat_count);
+        }
     }
 
     fn visit_expr(&mut self, expr: &'tcx Expr) {
@@ -814,7 +818,8 @@ impl<'tcx> ScopeTree {
 
     /// Checks whether the given scope contains a `yield`. If so,
     /// returns `Some((span, expr_count))` with the span of a yield we found and
-    /// the number of expressions appearing before the `yield` in the body.
+    /// the number of expressions and patterns appearing before the `yield` in the body + 1.
+    /// If there a are multiple yields in a scope, the one with the highest number is returned.
     pub fn yield_in_scope(&self, scope: Scope) -> Option<(Span, usize)> {
         self.yield_in_scope.get(&scope).cloned()
     }

src/librustc/traits/error_reporting.rs

@@ -1261,6 +1261,10 @@ impl<'a, 'gcx, 'tcx> InferCtxt<'a, 'gcx, 'tcx> {
                 err.note("the return type of a function must have a \
                           statically known size");
             }
+            ObligationCauseCode::SizedYieldType => {
+                err.note("the yield type of a generator must have a \
+                          statically known size");
+            }
             ObligationCauseCode::AssignmentLhsSized => {
                 err.note("the left-hand-side of an assignment must have a statically known size");
             }

src/librustc/traits/mod.rs

@@ -138,6 +138,8 @@ pub enum ObligationCauseCode<'tcx> {
     VariableType(ast::NodeId),
     /// Return type must be Sized
     SizedReturnType,
+    /// Yield type must be Sized
+    SizedYieldType,
     /// [T,..n] --> T must be Copy
     RepeatVec,
 

src/librustc/traits/structural_impls.rs

@@ -209,6 +209,7 @@ impl<'a, 'tcx> Lift<'tcx> for traits::ObligationCauseCode<'a> {
             super::VariableType(id) => Some(super::VariableType(id)),
             super::ReturnType(id) => Some(super::ReturnType(id)),
             super::SizedReturnType => Some(super::SizedReturnType),
+            super::SizedYieldType => Some(super::SizedYieldType),
             super::RepeatVec => Some(super::RepeatVec),
             super::FieldSized(item) => Some(super::FieldSized(item)),
             super::ConstSized => Some(super::ConstSized),
@@ -526,6 +527,7 @@ impl<'tcx> TypeFoldable<'tcx> for traits::ObligationCauseCode<'tcx> {
             super::VariableType(_) |
             super::ReturnType(_) |
             super::SizedReturnType |
+            super::SizedYieldType |
             super::ReturnNoExpression |
             super::RepeatVec |
             super::FieldSized(_) |
@@ -574,6 +576,7 @@ impl<'tcx> TypeFoldable<'tcx> for traits::ObligationCauseCode<'tcx> {
             super::VariableType(_) |
             super::ReturnType(_) |
             super::SizedReturnType |
+            super::SizedYieldType |
             super::ReturnNoExpression |
             super::RepeatVec |
             super::FieldSized(_) |

src/librustc/ty/sty.rs

@@ -390,14 +390,21 @@ impl<'a, 'gcx, 'tcx> ClosureSubsts<'tcx> {
         state.map(move |d| d.ty.subst(tcx, self.substs))
     }
 
+    /// This is the types of the fields of a generate which
+    /// is available before the generator transformation.
+    /// It includes the upvars and the state discriminant which is u32.
+    pub fn pre_transforms_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
+        impl Iterator<Item=Ty<'tcx>> + 'a
+    {
+        self.upvar_tys(def_id, tcx).chain(iter::once(tcx.types.u32))
+    }
+
     /// This is the types of all the fields stored in a generator.
     /// It includes the upvars, state types and the state discriminant which is u32.
     pub fn field_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
         impl Iterator<Item=Ty<'tcx>> + 'a
     {
-        let upvars = self.upvar_tys(def_id, tcx);
-        let state = self.state_tys(def_id, tcx);
-        upvars.chain(iter::once(tcx.types.u32)).chain(state)
+        self.pre_transforms_tys(def_id, tcx).chain(self.state_tys(def_id, tcx))
     }
 }
 

src/librustc_mir/borrow_check/error_reporting.rs

@@ -746,12 +746,12 @@ impl<'cx, 'gcx, 'tcx> MirBorrowckCtxt<'cx, 'gcx, 'tcx> {
                     self.describe_field_from_ty(&tnm.ty, field)
                 }
                 ty::TyArray(ty, _) | ty::TySlice(ty) => self.describe_field_from_ty(&ty, field),
-                ty::TyClosure(closure_def_id, _) => {
+                ty::TyClosure(def_id, _) | ty::TyGenerator(def_id, _, _) => {
                     // Convert the def-id into a node-id. node-ids are only valid for
                     // the local code in the current crate, so this returns an `Option` in case
                     // the closure comes from another crate. But in that case we wouldn't
                     // be borrowck'ing it, so we can just unwrap:
-                    let node_id = self.tcx.hir.as_local_node_id(closure_def_id).unwrap();
+                    let node_id = self.tcx.hir.as_local_node_id(def_id).unwrap();
                     let freevar = self.tcx.with_freevars(node_id, |fv| fv[field.index()]);
 
                     self.tcx.hir.name(freevar.var_id()).to_string()

src/librustc_mir/borrow_check/nll/type_check/mod.rs

@@ -533,15 +533,17 @@ impl<'a, 'b, 'gcx, 'tcx> TypeVerifier<'a, 'b, 'gcx, 'tcx> {
                     }
                 }
                 ty::TyGenerator(def_id, substs, _) => {
-                    // Try upvars first. `field_tys` requires final optimized MIR.
-                    if let Some(ty) = substs.upvar_tys(def_id, tcx).nth(field.index()) {
+                    // Try pre-transform fields first (upvars and current state)
+                    if let Some(ty) = substs.pre_transforms_tys(def_id, tcx).nth(field.index()) {
                         return Ok(ty);
                     }
 
+                    // Then try `field_tys` which contains all the fields, but it
+                    // requires the final optimized MIR.
                     return match substs.field_tys(def_id, tcx).nth(field.index()) {
                         Some(ty) => Ok(ty),
                         None => Err(FieldAccessError::OutOfRange {
-                            field_count: substs.field_tys(def_id, tcx).count() + 1,
+                            field_count: substs.field_tys(def_id, tcx).count(),
                         }),
                     };
                 }
@@ -1233,13 +1235,16 @@ impl<'a, 'gcx, 'tcx> TypeChecker<'a, 'gcx, 'tcx> {
                 }
             }
             AggregateKind::Generator(def_id, substs, _) => {
-                if let Some(ty) = substs.upvar_tys(def_id, tcx).nth(field_index) {
+                // Try pre-transform fields first (upvars and current state)
+                if let Some(ty) = substs.pre_transforms_tys(def_id, tcx).nth(field_index) {
                     Ok(ty)
                 } else {
+                    // Then try `field_tys` which contains all the fields, but it
+                    // requires the final optimized MIR.
                     match substs.field_tys(def_id, tcx).nth(field_index) {
                         Some(ty) => Ok(ty),
                         None => Err(FieldAccessError::OutOfRange {
-                            field_count: substs.field_tys(def_id, tcx).count() + 1,
+                            field_count: substs.field_tys(def_id, tcx).count(),
                         }),
                     }
                 }

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);
+            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);
+            }
+
+            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());
 
-            let mut live_locals = storage_liveness;
-
             // Mark locals without storage statements as always having live storage
-            live_locals.union(&ignored.0);
+            storage_liveness.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);
-            }
+            // 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]);
 
-            // 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 live_locals = storage_liveness;
 
             // Add the locals life at this suspension point to the set of locals which live across
             // any suspension points