Add a resume type parameter to Generator

Author: jonas-schievink

src/doc/unstable-book/src/language-features/generators.md

@@ -37,11 +37,11 @@ fn main() {
         return "foo"
     };
 
-    match Pin::new(&mut generator).resume() {
+    match Pin::new(&mut generator).resume(()) {
         GeneratorState::Yielded(1) => {}
         _ => panic!("unexpected value from resume"),
     }
-    match Pin::new(&mut generator).resume() {
+    match Pin::new(&mut generator).resume(()) {
         GeneratorState::Complete("foo") => {}
         _ => panic!("unexpected value from resume"),
     }
@@ -71,9 +71,9 @@ fn main() {
     };
 
     println!("1");
-    Pin::new(&mut generator).resume();
+    Pin::new(&mut generator).resume(());
     println!("3");
-    Pin::new(&mut generator).resume();
+    Pin::new(&mut generator).resume(());
     println!("5");
 }
 ```
@@ -92,10 +92,10 @@ The `Generator` trait in `std::ops` currently looks like:
 # use std::ops::GeneratorState;
 # use std::pin::Pin;
 
-pub trait Generator {
+pub trait Generator<R = ()> {
     type Yield;
     type Return;
-    fn resume(self: Pin<&mut Self>) -> GeneratorState<Self::Yield, Self::Return>;
+    fn resume(self: Pin<&mut Self>, resume: R) -> GeneratorState<Self::Yield, Self::Return>;
 }
 ```
 
@@ -152,10 +152,6 @@ closure-like semantics. Namely:
 * Whenever a generator is dropped it will drop all captured environment
   variables.
 
-Note that unlike closures, generators at this time cannot take any arguments.
-That is, generators must always look like `|| { ... }`. This restriction may be
-lifted at a future date, the design is ongoing!
-
 ### Generators as state machines
 
 In the compiler, generators are currently compiled as state machines. Each
@@ -179,8 +175,8 @@ fn main() {
         return ret
     };
 
-    Pin::new(&mut generator).resume();
-    Pin::new(&mut generator).resume();
+    Pin::new(&mut generator).resume(());
+    Pin::new(&mut generator).resume(());
 }
 ```
 
@@ -205,7 +201,7 @@ fn main() {
             type Yield = i32;
             type Return = &'static str;
 
-            fn resume(mut self: Pin<&mut Self>) -> GeneratorState<i32, &'static str> {
+            fn resume(mut self: Pin<&mut Self>, resume: ()) -> GeneratorState<i32, &'static str> {
                 use std::mem;
                 match mem::replace(&mut *self, __Generator::Done) {
                     __Generator::Start(s) => {
@@ -228,8 +224,8 @@ fn main() {
         __Generator::Start(ret)
     };
 
-    Pin::new(&mut generator).resume();
-    Pin::new(&mut generator).resume();
+    Pin::new(&mut generator).resume(());
+    Pin::new(&mut generator).resume(());
 }
 ```
 

src/liballoc/boxed.rs

@@ -1103,6 +1103,7 @@ impl<T: ?Sized> AsMut<T> for Box<T> {
 #[stable(feature = "pin", since = "1.33.0")]
 impl<T: ?Sized> Unpin for Box<T> {}
 
+#[cfg(bootstrap)]
 #[unstable(feature = "generator_trait", issue = "43122")]
 impl<G: ?Sized + Generator + Unpin> Generator for Box<G> {
     type Yield = G::Yield;
@@ -1113,6 +1114,7 @@ impl<G: ?Sized + Generator + Unpin> Generator for Box<G> {
     }
 }
 
+#[cfg(bootstrap)]
 #[unstable(feature = "generator_trait", issue = "43122")]
 impl<G: ?Sized + Generator> Generator for Pin<Box<G>> {
     type Yield = G::Yield;
@@ -1123,6 +1125,28 @@ impl<G: ?Sized + Generator> Generator for Pin<Box<G>> {
     }
 }
 
+#[cfg(not(bootstrap))]
+#[unstable(feature = "generator_trait", issue = "43122")]
+impl<G: ?Sized + Generator<R> + Unpin, R> Generator<R> for Box<G> {
+    type Yield = G::Yield;
+    type Return = G::Return;
+
+    fn resume(mut self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
+        G::resume(Pin::new(&mut *self), arg)
+    }
+}
+
+#[cfg(not(bootstrap))]
+#[unstable(feature = "generator_trait", issue = "43122")]
+impl<G: ?Sized + Generator<R>, R> Generator<R> for Pin<Box<G>> {
+    type Yield = G::Yield;
+    type Return = G::Return;
+
+    fn resume(mut self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
+        G::resume((*self).as_mut(), arg)
+    }
+}
+
 #[stable(feature = "futures_api", since = "1.36.0")]
 impl<F: ?Sized + Future + Unpin> Future for Box<F> {
     type Output = F::Output;

src/libcore/ops/generator.rs

@@ -50,11 +50,11 @@ pub enum GeneratorState<Y, R> {
 ///         return "foo"
 ///     };
 ///
-///     match Pin::new(&mut generator).resume() {
+///     match Pin::new(&mut generator).resume(()) {
 ///         GeneratorState::Yielded(1) => {}
 ///         _ => panic!("unexpected return from resume"),
 ///     }
-///     match Pin::new(&mut generator).resume() {
+///     match Pin::new(&mut generator).resume(()) {
 ///         GeneratorState::Complete("foo") => {}
 ///         _ => panic!("unexpected return from resume"),
 ///     }
@@ -67,7 +67,7 @@ pub enum GeneratorState<Y, R> {
 #[lang = "generator"]
 #[unstable(feature = "generator_trait", issue = "43122")]
 #[fundamental]
-pub trait Generator {
+pub trait Generator<#[cfg(not(bootstrap))] R = ()> {
     /// The type of value this generator yields.
     ///
     /// This associated type corresponds to the `yield` expression and the
@@ -110,9 +110,13 @@ pub trait Generator {
     /// been returned previously. While generator literals in the language are
     /// guaranteed to panic on resuming after `Complete`, this is not guaranteed
     /// for all implementations of the `Generator` trait.
-    fn resume(self: Pin<&mut Self>) -> GeneratorState<Self::Yield, Self::Return>;
+    fn resume(
+        self: Pin<&mut Self>,
+        #[cfg(not(bootstrap))] arg: R,
+    ) -> GeneratorState<Self::Yield, Self::Return>;
 }
 
+#[cfg(bootstrap)]
 #[unstable(feature = "generator_trait", issue = "43122")]
 impl<G: ?Sized + Generator> Generator for Pin<&mut G> {
     type Yield = G::Yield;
@@ -123,6 +127,7 @@ impl<G: ?Sized + Generator> Generator for Pin<&mut G> {
     }
 }
 
+#[cfg(bootstrap)]
 #[unstable(feature = "generator_trait", issue = "43122")]
 impl<G: ?Sized + Generator + Unpin> Generator for &mut G {
     type Yield = G::Yield;
@@ -132,3 +137,25 @@ impl<G: ?Sized + Generator + Unpin> Generator for &mut G {
         G::resume(Pin::new(&mut *self))
     }
 }
+
+#[cfg(not(bootstrap))]
+#[unstable(feature = "generator_trait", issue = "43122")]
+impl<G: ?Sized + Generator<R>, R> Generator<R> for Pin<&mut G> {
+    type Yield = G::Yield;
+    type Return = G::Return;
+
+    fn resume(mut self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
+        G::resume((*self).as_mut(), arg)
+    }
+}
+
+#[cfg(not(bootstrap))]
+#[unstable(feature = "generator_trait", issue = "43122")]
+impl<G: ?Sized + Generator<R> + Unpin, R> Generator<R> for &mut G {
+    type Yield = G::Yield;
+    type Return = G::Return;
+
+    fn resume(mut self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return> {
+        G::resume(Pin::new(&mut *self), arg)
+    }
+}

src/librustc/traits/util.rs

@@ -643,8 +643,10 @@ pub fn generator_trait_ref_and_outputs(
     self_ty: Ty<'tcx>,
     sig: ty::PolyGenSig<'tcx>,
 ) -> ty::Binder<(ty::TraitRef<'tcx>, Ty<'tcx>, Ty<'tcx>)> {
-    let trait_ref =
-        ty::TraitRef { def_id: fn_trait_def_id, substs: tcx.mk_substs_trait(self_ty, &[]) };
+    let trait_ref = ty::TraitRef {
+        def_id: fn_trait_def_id,
+        substs: tcx.mk_substs_trait(self_ty, &[tcx.mk_unit().into()]),
+    };
     ty::Binder::bind((trait_ref, sig.skip_binder().yield_ty, sig.skip_binder().return_ty))
 }
 

src/librustc/ty/layout.rs

@@ -2350,8 +2350,8 @@ impl<'tcx> ty::Instance<'tcx> {
                     ]);
                     let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
 
-                    tcx.mk_fn_sig(iter::once(env_ty),
-                        ret_ty,
+                    tcx.mk_fn_sig([env_ty, tcx.mk_unit()].iter(),
+                        &ret_ty,
                         false,
                         hir::Unsafety::Normal,
                         rustc_target::spec::abi::Abi::Rust

src/librustc_error_codes/error_codes/E0626.md

@@ -12,7 +12,7 @@ let mut b = || {
     yield (); // ...is still in scope here, when the yield occurs.
     println!("{}", a);
 };
-Pin::new(&mut b).resume();
+Pin::new(&mut b).resume(());
 ```
 
 At present, it is not permitted to have a yield that occurs while a
@@ -31,7 +31,7 @@ let mut b = || {
     yield ();
     println!("{}", a);
 };
-Pin::new(&mut b).resume();
+Pin::new(&mut b).resume(());
 ```
 
 This is a very simple case, of course. In more complex cases, we may
@@ -50,7 +50,7 @@ let mut b = || {
     yield x; // ...when this yield occurs.
   }
 };
-Pin::new(&mut b).resume();
+Pin::new(&mut b).resume(());
 ```
 
 Such cases can sometimes be resolved by iterating "by value" (or using
@@ -66,7 +66,7 @@ let mut b = || {
     yield x; // <-- Now yield is OK.
   }
 };
-Pin::new(&mut b).resume();
+Pin::new(&mut b).resume(());
 ```
 
 If taking ownership is not an option, using indices can work too:
@@ -83,7 +83,7 @@ let mut b = || {
     yield x; // <-- Now yield is OK.
   }
 };
-Pin::new(&mut b).resume();
+Pin::new(&mut b).resume(());
 
 // (*) -- Unfortunately, these temporaries are currently required.
 // See <https://github.com/rust-lang/rust/issues/43122>.

src/librustc_mir/dataflow/impls/storage_liveness.rs

@@ -31,10 +31,12 @@ impl<'a, 'tcx> BitDenotation<'tcx> for MaybeStorageLive<'a, 'tcx> {
         self.body.local_decls.len()
     }
 
-    fn start_block_effect(&self, _on_entry: &mut BitSet<Local>) {
-        // Nothing is live on function entry (generators only have a self
-        // argument, and we don't care about that)
-        assert_eq!(1, self.body.arg_count);
+    fn start_block_effect(&self, on_entry: &mut BitSet<Local>) {
+        // The resume argument is live on function entry (we don't care about
+        // the `self` argument)
+        for arg in self.body.args_iter().skip(1) {
+            on_entry.insert(arg);
+        }
     }
 
     fn statement_effect(&self, trans: &mut GenKillSet<Local>, loc: Location) {
@@ -100,10 +102,12 @@ impl<'mir, 'tcx> BitDenotation<'tcx> for RequiresStorage<'mir, 'tcx> {
         self.body.local_decls.len()
     }
 
-    fn start_block_effect(&self, _sets: &mut BitSet<Local>) {
-        // Nothing is live on function entry (generators only have a self
-        // argument, and we don't care about that)
-        assert_eq!(1, self.body.arg_count);
+    fn start_block_effect(&self, on_entry: &mut BitSet<Local>) {
+        // The resume argument is live on function entry (we don't care about
+        // the `self` argument)
+        for arg in self.body.args_iter().skip(1) {
+            on_entry.insert(arg);
+        }
     }
 
     fn before_statement_effect(&self, sets: &mut GenKillSet<Self::Idx>, loc: Location) {

src/librustc_mir/transform/generator.rs

@@ -885,6 +885,7 @@ fn create_generator_drop_shim<'tcx>(
     drop_clean: BasicBlock,
 ) -> BodyAndCache<'tcx> {
     let mut body = body.clone();
+    body.arg_count = 1; // make sure the resume argument is not included here
 
     let source_info = source_info(&body);
 
@@ -1164,7 +1165,7 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
 
         // Update our MIR struct to reflect the changed we've made
         body.yield_ty = None;
-        body.arg_count = 1;
+        body.arg_count = 2; // self, resume arg
         body.spread_arg = None;
         body.generator_layout = Some(layout);
 

src/librustc_mir_build/build/mod.rs

@@ -75,13 +75,16 @@ fn mir_build(tcx: TyCtxt<'_>, def_id: DefId) -> BodyAndCache<'_> {
                     // HACK(eddyb) Avoid having RustCall on closures,
                     // as it adds unnecessary (and wrong) auto-tupling.
                     abi = Abi::Rust;
-                    Some(ArgInfo(liberated_closure_env_ty(tcx, id, body_id), None, None, None))
+                    vec![ArgInfo(liberated_closure_env_ty(tcx, id, body_id), None, None, None)]
                 }
                 ty::Generator(..) => {
                     let gen_ty = tcx.body_tables(body_id).node_type(id);
-                    Some(ArgInfo(gen_ty, None, None, None))
+                    vec![
+                        ArgInfo(gen_ty, None, None, None),
+                        ArgInfo(tcx.mk_unit(), None, None, None),
+                    ]
                 }
-                _ => None,
+                _ => vec![],
             };
 
             let safety = match fn_sig.unsafety {

src/libstd/future.rs

@@ -40,7 +40,10 @@ impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
         // Safe because we're !Unpin + !Drop mapping to a ?Unpin value
         let gen = unsafe { Pin::map_unchecked_mut(self, |s| &mut s.0) };
         let _guard = unsafe { set_task_context(cx) };
-        match gen.resume() {
+        match gen.resume(
+            #[cfg(not(bootstrap))]
+            (),
+        ) {
             GeneratorState::Yielded(()) => Poll::Pending,
             GeneratorState::Complete(x) => Poll::Ready(x),
         }

src/test/debuginfo/generator-locals.rs

@@ -78,9 +78,9 @@ fn main() {
         _zzz(); // #break
         a = c;
     };
-    Pin::new(&mut b).resume();
-    Pin::new(&mut b).resume();
-    Pin::new(&mut b).resume();
+    Pin::new(&mut b).resume(());
+    Pin::new(&mut b).resume(());
+    Pin::new(&mut b).resume(());
     _zzz(); // #break
 }
 

src/test/debuginfo/generator-objects.rs

@@ -57,11 +57,11 @@ fn main() {
         println!("{} {} {}", a, c, d);
     };
     _zzz(); // #break
-    Pin::new(&mut b).resume();
+    Pin::new(&mut b).resume(());
     _zzz(); // #break
-    Pin::new(&mut b).resume();
+    Pin::new(&mut b).resume(());
     _zzz(); // #break
-    Pin::new(&mut b).resume();
+    Pin::new(&mut b).resume(());
     _zzz(); // #break
 }
 

src/test/debuginfo/issue-57822.rs

@@ -45,7 +45,7 @@ fn main() {
         yield;
     };
     let mut b = move || {
-        Pin::new(&mut a).resume();
+        Pin::new(&mut a).resume(());
         yield;
     };
 

src/test/run-fail/generator-resume-after-panic.rs

@@ -16,7 +16,7 @@ fn main() {
         yield;
     };
     panic::catch_unwind(panic::AssertUnwindSafe(|| {
-        let x = Pin::new(&mut g).resume();
+        let x = Pin::new(&mut g).resume(());
     }));
-    Pin::new(&mut g).resume();
+    Pin::new(&mut g).resume(());
 }