Allow access of the state field before the generator transform. Fixes #47482, #46476

4 files changed, 38 insertions, 26 deletions

diff --git a/src/librustc/ty/sty.rs b/src/librustc/ty/sty.rs
index db7e4fe45ef..d20cf556b71 100644
--- a/src/librustc/ty/sty.rs
+++ b/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))
     }
 }
 
diff --git a/src/librustc_mir/borrow_check/nll/type_check/mod.rs b/src/librustc_mir/borrow_check/nll/type_check/mod.rs
index 9dcd4435580..dc302b6cc35 100644
--- a/src/librustc_mir/borrow_check/nll/type_check/mod.rs
+++ b/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(),
                         }),
                     }
                 }
diff --git a/src/test/ui/generator/generator-with-nll.stderr b/src/test/ui/generator/generator-with-nll.stderr
index 0a52a928f69..0f7d2e540d8 100644
--- a/src/test/ui/generator/generator-with-nll.stderr
+++ b/src/test/ui/generator/generator-with-nll.stderr
@@ -1,12 +1,3 @@
-error[E0626]: borrow may still be in use when generator yields (Mir)
-  --> $DIR/generator-with-nll.rs:20:17
-   |
-20 |         let b = &mut true; //~ ERROR borrow may still be in use when generator yields (Ast)
-   |                 ^^^^^^^^^
-21 |         //~^ borrow may still be in use when generator yields (Mir)
-22 |         yield ();
-   |         -------- possible yield occurs here
-
 error[E0626]: borrow may still be in use when generator yields (Ast)
   --> $DIR/generator-with-nll.rs:19:23
    |
@@ -25,5 +16,14 @@ error[E0626]: borrow may still be in use when generator yields (Ast)
 22 |         yield ();
    |         -------- possible yield occurs here
 
+error[E0626]: borrow may still be in use when generator yields (Mir)
+  --> $DIR/generator-with-nll.rs:20:17
+   |
+20 |         let b = &mut true; //~ ERROR borrow may still be in use when generator yields (Ast)
+   |                 ^^^^^^^^^
+21 |         //~^ borrow may still be in use when generator yields (Mir)
+22 |         yield ();
+   |         -------- possible yield occurs here
+
 error: aborting due to 3 previous errors
 
diff --git a/src/test/ui/generator/yield-while-local-borrowed.stderr b/src/test/ui/generator/yield-while-local-borrowed.stderr
index 7961dd97441..114fe8ffcab 100644
--- a/src/test/ui/generator/yield-while-local-borrowed.stderr
+++ b/src/test/ui/generator/yield-while-local-borrowed.stderr
@@ -1,12 +1,3 @@
-error[E0626]: borrow may still be in use when generator yields (Mir)
-  --> $DIR/yield-while-local-borrowed.rs:24:17
-   |
-24 |         let a = &mut 3;
-   |                 ^^^^^^
-...
-27 |         yield();
-   |         ------- possible yield occurs here
-
 error[E0626]: borrow may still be in use when generator yields (Ast)
   --> $DIR/yield-while-local-borrowed.rs:24:22
    |
@@ -26,6 +17,15 @@ error[E0626]: borrow may still be in use when generator yields (Ast)
    |             ------- possible yield occurs here
 
 error[E0626]: borrow may still be in use when generator yields (Mir)
+  --> $DIR/yield-while-local-borrowed.rs:24:17
+   |
+24 |         let a = &mut 3;
+   |                 ^^^^^^
+...
+27 |         yield();
+   |         ------- possible yield occurs here
+
+error[E0626]: borrow may still be in use when generator yields (Mir)
   --> $DIR/yield-while-local-borrowed.rs:52:21
    |
 52 |             let b = &a;