Rollup merge of #68524 - jonas-schievink:generator-resume-arguments, r=Zoxc

Generator Resume Arguments

cc #43122 and #56974

Blockers:
* [x] Fix miscompilation when resume argument is live across a yield point (#68524 (comment))
* [x] Fix 10% compile time regression in `await-call-tree` benchmarks (#68524 (comment))
  * [x] Fix remaining 1-3% regression (#68524 (comment)) - resolved (#68524 (comment))
* [x] Make dropck rules account for resume arguments (#68524 (comment))

Follow-up work:
* Change async/await desugaring to make use of this feature
* Rewrite [`box_region.rs`](https://github.com/rust-lang/rust/blob/3d8778d767f0dde6fe2bc9459f21ead8e124d8cb/src/librustc_data_structures/box_region.rs) to use resume arguments (this shows up in profiles too)

Author: Dylan-DPC

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

@@ -1104,6 +1104,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;
@@ -1114,6 +1115,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;
@@ -1124,6 +1126,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/infer/opaque_types/mod.rs

@@ -744,6 +744,7 @@ where
 
                 substs.as_generator().return_ty(def_id, self.tcx).visit_with(self);
                 substs.as_generator().yield_ty(def_id, self.tcx).visit_with(self);
+                substs.as_generator().resume_ty(def_id, self.tcx).visit_with(self);
             }
             _ => {
                 ty.super_visit_with(self);

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));

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);
                     }
 

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, &[sig.skip_binder().resume_ty.into()]),
+    };
     ty::Binder::bind((trait_ref, sig.skip_binder().yield_ty, sig.skip_binder().return_ty))
 }
 

src/librustc/ty/layout.rs

@@ -2350,8 +2350,9 @@ 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, sig.resume_ty].iter(),
+                        &ret_ty,
                         false,
                         hir::Unsafety::Normal,
                         rustc_target::spec::abi::Abi::Rust

src/librustc/ty/structural_impls.rs

@@ -598,8 +598,8 @@ impl<'a, 'tcx> Lift<'tcx> for ty::adjustment::AutoBorrow<'a> {
 impl<'a, 'tcx> Lift<'tcx> for ty::GenSig<'a> {
     type Lifted = ty::GenSig<'tcx>;
     fn lift_to_tcx(&self, tcx: TyCtxt<'tcx>) -> Option<Self::Lifted> {
-        tcx.lift(&(self.yield_ty, self.return_ty))
-            .map(|(yield_ty, return_ty)| ty::GenSig { yield_ty, return_ty })
+        tcx.lift(&(self.resume_ty, self.yield_ty, self.return_ty))
+            .map(|(resume_ty, yield_ty, return_ty)| ty::GenSig { resume_ty, yield_ty, return_ty })
     }
 }
 

src/librustc/ty/sty.rs

@@ -346,9 +346,17 @@ static_assert_size!(TyKind<'_>, 24);
 /// ## Generators
 ///
 /// Generators are handled similarly in `GeneratorSubsts`.  The set of
-/// type parameters is similar, but the role of CK and CS are
-/// different. CK represents the "yield type" and CS represents the
-/// "return type" of the generator.
+/// type parameters is similar, but `CK` and `CS` are replaced by the
+/// following type parameters:
+///
+/// * `GS`: The generator's "resume type", which is the type of the
+///   argument passed to `resume`, and the type of `yield` expressions
+///   inside the generator.
+/// * `GY`: The "yield type", which is the type of values passed to
+///   `yield` inside the generator.
+/// * `GR`: The "return type", which is the type of value returned upon
+///   completion of the generator.
+/// * `GW`: The "generator witness".
 #[derive(Copy, Clone, Debug, TypeFoldable)]
 pub struct ClosureSubsts<'tcx> {
     /// Lifetime and type parameters from the enclosing function,
@@ -442,6 +450,7 @@ pub struct GeneratorSubsts<'tcx> {
 }
 
 struct SplitGeneratorSubsts<'tcx> {
+    resume_ty: Ty<'tcx>,
     yield_ty: Ty<'tcx>,
     return_ty: Ty<'tcx>,
     witness: Ty<'tcx>,
@@ -453,10 +462,11 @@ impl<'tcx> GeneratorSubsts<'tcx> {
         let generics = tcx.generics_of(def_id);
         let parent_len = generics.parent_count;
         SplitGeneratorSubsts {
-            yield_ty: self.substs.type_at(parent_len),
-            return_ty: self.substs.type_at(parent_len + 1),
-            witness: self.substs.type_at(parent_len + 2),
-            upvar_kinds: &self.substs[parent_len + 3..],
+            resume_ty: self.substs.type_at(parent_len),
+            yield_ty: self.substs.type_at(parent_len + 1),
+            return_ty: self.substs.type_at(parent_len + 2),
+            witness: self.substs.type_at(parent_len + 3),
+            upvar_kinds: &self.substs[parent_len + 4..],
         }
     }
 
@@ -485,6 +495,11 @@ impl<'tcx> GeneratorSubsts<'tcx> {
         })
     }
 
+    /// Returns the type representing the resume type of the generator.
+    pub fn resume_ty(self, def_id: DefId, tcx: TyCtxt<'_>) -> Ty<'tcx> {
+        self.split(def_id, tcx).resume_ty
+    }
+
     /// Returns the type representing the yield type of the generator.
     pub fn yield_ty(self, def_id: DefId, tcx: TyCtxt<'_>) -> Ty<'tcx> {
         self.split(def_id, tcx).yield_ty
@@ -505,10 +520,14 @@ impl<'tcx> GeneratorSubsts<'tcx> {
         ty::Binder::dummy(self.sig(def_id, tcx))
     }
 
-    /// Returns the "generator signature", which consists of its yield
+    /// Returns the "generator signature", which consists of its resume, yield
     /// and return types.
     pub fn sig(self, def_id: DefId, tcx: TyCtxt<'_>) -> GenSig<'tcx> {
-        ty::GenSig { yield_ty: self.yield_ty(def_id, tcx), return_ty: self.return_ty(def_id, tcx) }
+        ty::GenSig {
+            resume_ty: self.resume_ty(def_id, tcx),
+            yield_ty: self.yield_ty(def_id, tcx),
+            return_ty: self.return_ty(def_id, tcx),
+        }
     }
 }
 
@@ -1072,13 +1091,17 @@ impl<'tcx> ProjectionTy<'tcx> {
 
 #[derive(Clone, Debug, TypeFoldable)]
 pub struct GenSig<'tcx> {
+    pub resume_ty: Ty<'tcx>,
     pub yield_ty: Ty<'tcx>,
     pub return_ty: Ty<'tcx>,
 }
 
 pub type PolyGenSig<'tcx> = Binder<GenSig<'tcx>>;
 
 impl<'tcx> PolyGenSig<'tcx> {
+    pub fn resume_ty(&self) -> ty::Binder<Ty<'tcx>> {
+        self.map_bound_ref(|sig| sig.resume_ty)
+    }
     pub fn yield_ty(&self) -> ty::Binder<Ty<'tcx>> {
         self.map_bound_ref(|sig| sig.yield_ty)
     }