Add resume arg place to `Yield` MIR terminator

6 files changed, 26 insertions, 7 deletions

diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index 4e9835a52b7..f6c7174649f 100644
--- a/src/librustc/mir/mod.rs
+++ b/src/librustc/mir/mod.rs
@@ -1120,6 +1120,8 @@ pub enum TerminatorKind<'tcx> {
         value: Operand<'tcx>,
         /// Where to resume to.
         resume: BasicBlock,
+        /// The place to store the resume argument in.
+        resume_arg: Place<'tcx>,
         /// Cleanup to be done if the generator is dropped at this suspend point.
         drop: Option<BasicBlock>,
     },
@@ -2645,9 +2647,12 @@ impl<'tcx> TypeFoldable<'tcx> for Terminator<'tcx> {
                 target,
                 unwind,
             },
-            Yield { ref value, resume, drop } => {
-                Yield { value: value.fold_with(folder), resume: resume, drop: drop }
-            }
+            Yield { ref value, resume, ref resume_arg, drop } => Yield {
+                value: value.fold_with(folder),
+                resume,
+                resume_arg: resume_arg.fold_with(folder),
+                drop,
+            },
             Call { ref func, ref args, ref destination, cleanup, from_hir_call } => {
                 let dest =
                     destination.as_ref().map(|&(ref loc, dest)| (loc.fold_with(folder), dest));
diff --git a/src/librustc/mir/visit.rs b/src/librustc/mir/visit.rs
index 4c5db1b07d2..2f094516a35 100644
--- a/src/librustc/mir/visit.rs
+++ b/src/librustc/mir/visit.rs
@@ -516,8 +516,14 @@ macro_rules! make_mir_visitor {
                     TerminatorKind::Yield {
                         value,
                         resume: _,
+                        resume_arg,
                         drop: _,
                     } => {
+                        self.visit_place(
+                            resume_arg,
+                            PlaceContext::MutatingUse(MutatingUseContext::Store),
+                            source_location,
+                        );
                         self.visit_operand(value, source_location);
                     }
 
diff --git a/src/librustc_mir/borrow_check/invalidation.rs b/src/librustc_mir/borrow_check/invalidation.rs
index bb56c11872a..392f164d314 100644
--- a/src/librustc_mir/borrow_check/invalidation.rs
+++ b/src/librustc_mir/borrow_check/invalidation.rs
@@ -159,7 +159,7 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                     self.consume_operand(location, index);
                 }
             }
-            TerminatorKind::Yield { ref value, resume, drop: _ } => {
+            TerminatorKind::Yield { ref value, resume, resume_arg, drop: _ } => {
                 self.consume_operand(location, value);
 
                 // Invalidate all borrows of local places
@@ -170,6 +170,8 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                         self.all_facts.invalidates.push((resume, i));
                     }
                 }
+
+                self.mutate_place(location, resume_arg, Deep, JustWrite);
             }
             TerminatorKind::Resume | TerminatorKind::Return | TerminatorKind::GeneratorDrop => {
                 // Invalidate all borrows of local places
diff --git a/src/librustc_mir/borrow_check/mod.rs b/src/librustc_mir/borrow_check/mod.rs
index 717359d75c3..e528159fcef 100644
--- a/src/librustc_mir/borrow_check/mod.rs
+++ b/src/librustc_mir/borrow_check/mod.rs
@@ -684,7 +684,7 @@ impl<'cx, 'tcx> DataflowResultsConsumer<'cx, 'tcx> for MirBorrowckCtxt<'cx, 'tcx
                 }
             }
 
-            TerminatorKind::Yield { ref value, resume: _, drop: _ } => {
+            TerminatorKind::Yield { ref value, resume: _, ref resume_arg, drop: _ } => {
                 self.consume_operand(loc, (value, span), flow_state);
 
                 if self.movable_generator {
@@ -697,6 +697,8 @@ impl<'cx, 'tcx> DataflowResultsConsumer<'cx, 'tcx> for MirBorrowckCtxt<'cx, 'tcx
                         }
                     });
                 }
+
+                self.mutate_place(loc, (resume_arg, span), Deep, JustWrite, flow_state);
             }
 
             TerminatorKind::Resume | TerminatorKind::Return | TerminatorKind::GeneratorDrop => {
diff --git a/src/librustc_mir/borrow_check/universal_regions.rs b/src/librustc_mir/borrow_check/universal_regions.rs
index 6e36508ed60..f6e3ca2f809 100644
--- a/src/librustc_mir/borrow_check/universal_regions.rs
+++ b/src/librustc_mir/borrow_check/universal_regions.rs
@@ -581,9 +581,11 @@ impl<'cx, 'tcx> UniversalRegionsBuilder<'cx, 'tcx> {
 
             DefiningTy::Generator(def_id, substs, movability) => {
                 assert_eq!(self.mir_def_id, def_id);
+                let resume_ty = substs.as_generator().resume_ty(def_id, tcx);
                 let output = substs.as_generator().return_ty(def_id, tcx);
                 let generator_ty = tcx.mk_generator(def_id, substs, movability);
-                let inputs_and_output = self.infcx.tcx.intern_type_list(&[generator_ty, output]);
+                let inputs_and_output =
+                    self.infcx.tcx.intern_type_list(&[generator_ty, resume_ty, output]);
                 ty::Binder::dummy(inputs_and_output)
             }
 
diff --git a/src/librustc_mir/dataflow/move_paths/builder.rs b/src/librustc_mir/dataflow/move_paths/builder.rs
index 62af196174f..6f8caca5e21 100644
--- a/src/librustc_mir/dataflow/move_paths/builder.rs
+++ b/src/librustc_mir/dataflow/move_paths/builder.rs
@@ -380,7 +380,9 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
                 self.gather_operand(discr);
             }
 
-            TerminatorKind::Yield { ref value, .. } => {
+            TerminatorKind::Yield { ref value, resume_arg: ref place, .. } => {
+                self.create_move_path(place);
+                self.gather_init(place.as_ref(), InitKind::Deep);
                 self.gather_operand(value);
             }