Rollup merge of #69033 - jonas-schievink:resume-with-context, r=tmandry

Use generator resume arguments in the async/await lowering

This removes the TLS requirement from async/await and enables it in `#![no_std]` crates.

Closes https://github.com/rust-lang/rust/issues/56974

I'm not confident the HIR lowering is completely correct, there seem to be quite a few undocumented invariants in there. The `async-std` and tokio test suites are passing with these changes though.

9 files changed, 204 insertions, 36 deletions

diff --git a/src/libcore/future/mod.rs b/src/libcore/future/mod.rs
index 89ea4713cfd..8dfda7a4a32 100644
--- a/src/libcore/future/mod.rs
+++ b/src/libcore/future/mod.rs
@@ -2,6 +2,84 @@
 
 //! Asynchronous values.
 
+#[cfg(not(bootstrap))]
+use crate::{
+    ops::{Generator, GeneratorState},
+    pin::Pin,
+    ptr::NonNull,
+    task::{Context, Poll},
+};
+
 mod future;
 #[stable(feature = "futures_api", since = "1.36.0")]
 pub use self::future::Future;
+
+/// This type is needed because:
+///
+/// a) Generators cannot implement `for<'a, 'b> Generator<&'a mut Context<'b>>`, so we need to pass
+///    a raw pointer (see https://github.com/rust-lang/rust/issues/68923).
+/// b) Raw pointers and `NonNull` aren't `Send` or `Sync`, so that would make every single future
+///    non-Send/Sync as well, and we don't want that.
+///
+/// It also simplifies the HIR lowering of `.await`.
+#[doc(hidden)]
+#[unstable(feature = "gen_future", issue = "50547")]
+#[cfg(not(bootstrap))]
+#[derive(Debug, Copy, Clone)]
+pub struct ResumeTy(NonNull<Context<'static>>);
+
+#[unstable(feature = "gen_future", issue = "50547")]
+#[cfg(not(bootstrap))]
+unsafe impl Send for ResumeTy {}
+
+#[unstable(feature = "gen_future", issue = "50547")]
+#[cfg(not(bootstrap))]
+unsafe impl Sync for ResumeTy {}
+
+/// Wrap a generator in a future.
+///
+/// This function returns a `GenFuture` underneath, but hides it in `impl Trait` to give
+/// better error messages (`impl Future` rather than `GenFuture<[closure.....]>`).
+// This is `const` to avoid extra errors after we recover from `const async fn`
+#[doc(hidden)]
+#[unstable(feature = "gen_future", issue = "50547")]
+#[cfg(not(bootstrap))]
+#[inline]
+pub const fn from_generator<T>(gen: T) -> impl Future<Output = T::Return>
+where
+    T: Generator<ResumeTy, Yield = ()>,
+{
+    struct GenFuture<T: Generator<ResumeTy, Yield = ()>>(T);
+
+    // We rely on the fact that async/await futures are immovable in order to create
+    // self-referential borrows in the underlying generator.
+    impl<T: Generator<ResumeTy, Yield = ()>> !Unpin for GenFuture<T> {}
+
+    impl<T: Generator<ResumeTy, Yield = ()>> Future for GenFuture<T> {
+        type Output = T::Return;
+        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            // Safety: Safe because we're !Unpin + !Drop, and this is just a field projection.
+            let gen = unsafe { Pin::map_unchecked_mut(self, |s| &mut s.0) };
+
+            // Resume the generator, turning the `&mut Context` into a `NonNull` raw pointer. The
+            // `.await` lowering will safely cast that back to a `&mut Context`.
+            match gen.resume(ResumeTy(NonNull::from(cx).cast::<Context<'static>>())) {
+                GeneratorState::Yielded(()) => Poll::Pending,
+                GeneratorState::Complete(x) => Poll::Ready(x),
+            }
+        }
+    }
+
+    GenFuture(gen)
+}
+
+#[doc(hidden)]
+#[unstable(feature = "gen_future", issue = "50547")]
+#[cfg(not(bootstrap))]
+#[inline]
+pub unsafe fn poll_with_context<F>(f: Pin<&mut F>, mut cx: ResumeTy) -> Poll<F::Output>
+where
+    F: Future,
+{
+    F::poll(f, cx.0.as_mut())
+}
diff --git a/src/librustc_ast_lowering/expr.rs b/src/librustc_ast_lowering/expr.rs
index a4cbae51966..a212f0a7107 100644
--- a/src/librustc_ast_lowering/expr.rs
+++ b/src/librustc_ast_lowering/expr.rs
@@ -470,6 +470,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
         }
     }
 
+    /// Lower an `async` construct to a generator that is then wrapped so it implements `Future`.
+    ///
+    /// This results in:
+    ///
+    /// ```text
+    /// std::future::from_generator(static move? |_task_context| -> <ret_ty> {
+    ///     <body>
+    /// })
+    /// ```
     pub(super) fn make_async_expr(
         &mut self,
         capture_clause: CaptureBy,
@@ -480,17 +489,42 @@ impl<'hir> LoweringContext<'_, 'hir> {
         body: impl FnOnce(&mut Self) -> hir::Expr<'hir>,
     ) -> hir::ExprKind<'hir> {
         let output = match ret_ty {
-            Some(ty) => FnRetTy::Ty(ty),
-            None => FnRetTy::Default(span),
+            Some(ty) => hir::FnRetTy::Return(self.lower_ty(&ty, ImplTraitContext::disallowed())),
+            None => hir::FnRetTy::DefaultReturn(span),
         };
-        let ast_decl = FnDecl { inputs: vec![], output };
-        let decl = self.lower_fn_decl(&ast_decl, None, /* impl trait allowed */ false, None);
-        let body_id = self.lower_fn_body(&ast_decl, |this| {
+
+        // Resume argument type. We let the compiler infer this to simplify the lowering. It is
+        // fully constrained by `future::from_generator`.
+        let input_ty = hir::Ty { hir_id: self.next_id(), kind: hir::TyKind::Infer, span };
+
+        // The closure/generator `FnDecl` takes a single (resume) argument of type `input_ty`.
+        let decl = self.arena.alloc(hir::FnDecl {
+            inputs: arena_vec![self; input_ty],
+            output,
+            c_variadic: false,
+            implicit_self: hir::ImplicitSelfKind::None,
+        });
+
+        // Lower the argument pattern/ident. The ident is used again in the `.await` lowering.
+        let (pat, task_context_hid) = self.pat_ident_binding_mode(
+            span,
+            Ident::with_dummy_span(sym::_task_context),
+            hir::BindingAnnotation::Mutable,
+        );
+        let param = hir::Param { attrs: &[], hir_id: self.next_id(), pat, span };
+        let params = arena_vec![self; param];
+
+        let body_id = self.lower_body(move |this| {
             this.generator_kind = Some(hir::GeneratorKind::Async(async_gen_kind));
-            body(this)
+
+            let old_ctx = this.task_context;
+            this.task_context = Some(task_context_hid);
+            let res = body(this);
+            this.task_context = old_ctx;
+            (params, res)
         });
 
-        // `static || -> <ret_ty> { body }`:
+        // `static |_task_context| -> <ret_ty> { body }`:
         let generator_kind = hir::ExprKind::Closure(
             capture_clause,
             decl,
@@ -523,13 +557,14 @@ impl<'hir> LoweringContext<'_, 'hir> {
     /// ```rust
     /// match <expr> {
     ///     mut pinned => loop {
-    ///         match ::std::future::poll_with_tls_context(unsafe {
-    ///             <::std::pin::Pin>::new_unchecked(&mut pinned)
-    ///         }) {
+    ///         match unsafe { ::std::future::poll_with_context(
+    ///             <::std::pin::Pin>::new_unchecked(&mut pinned),
+    ///             task_context,
+    ///         ) } {
     ///             ::std::task::Poll::Ready(result) => break result,
     ///             ::std::task::Poll::Pending => {}
     ///         }
-    ///         yield ();
+    ///         task_context = yield ();
     ///     }
     /// }
     /// ```
@@ -561,12 +596,23 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let (pinned_pat, pinned_pat_hid) =
             self.pat_ident_binding_mode(span, pinned_ident, hir::BindingAnnotation::Mutable);
 
-        // ::std::future::poll_with_tls_context(unsafe {
-        //     ::std::pin::Pin::new_unchecked(&mut pinned)
-        // })`
+        let task_context_ident = Ident::with_dummy_span(sym::_task_context);
+
+        // unsafe {
+        //     ::std::future::poll_with_context(
+        //         ::std::pin::Pin::new_unchecked(&mut pinned),
+        //         task_context,
+        //     )
+        // }
         let poll_expr = {
             let pinned = self.expr_ident(span, pinned_ident, pinned_pat_hid);
             let ref_mut_pinned = self.expr_mut_addr_of(span, pinned);
+            let task_context = if let Some(task_context_hid) = self.task_context {
+                self.expr_ident_mut(span, task_context_ident, task_context_hid)
+            } else {
+                // Use of `await` outside of an async context, we cannot use `task_context` here.
+                self.expr_err(span)
+            };
             let pin_ty_id = self.next_id();
             let new_unchecked_expr_kind = self.expr_call_std_assoc_fn(
                 pin_ty_id,
@@ -575,14 +621,13 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 "new_unchecked",
                 arena_vec![self; ref_mut_pinned],
             );
-            let new_unchecked =
-                self.arena.alloc(self.expr(span, new_unchecked_expr_kind, ThinVec::new()));
-            let unsafe_expr = self.expr_unsafe(new_unchecked);
-            self.expr_call_std_path(
+            let new_unchecked = self.expr(span, new_unchecked_expr_kind, ThinVec::new());
+            let call = self.expr_call_std_path(
                 gen_future_span,
-                &[sym::future, sym::poll_with_tls_context],
-                arena_vec![self; unsafe_expr],
-            )
+                &[sym::future, sym::poll_with_context],
+                arena_vec![self; new_unchecked, task_context],
+            );
+            self.arena.alloc(self.expr_unsafe(call))
         };
 
         // `::std::task::Poll::Ready(result) => break result`
@@ -622,6 +667,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             self.stmt_expr(span, match_expr)
         };
 
+        // task_context = yield ();
         let yield_stmt = {
             let unit = self.expr_unit(span);
             let yield_expr = self.expr(
@@ -629,7 +675,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 hir::ExprKind::Yield(unit, hir::YieldSource::Await),
                 ThinVec::new(),
             );
-            self.stmt_expr(span, yield_expr)
+            let yield_expr = self.arena.alloc(yield_expr);
+
+            if let Some(task_context_hid) = self.task_context {
+                let lhs = self.expr_ident(span, task_context_ident, task_context_hid);
+                let assign =
+                    self.expr(span, hir::ExprKind::Assign(lhs, yield_expr, span), AttrVec::new());
+                self.stmt_expr(span, assign)
+            } else {
+                // Use of `await` outside of an async context. Return `yield_expr` so that we can
+                // proceed with type checking.
+                self.stmt(span, hir::StmtKind::Semi(yield_expr))
+            }
         };
 
         let loop_block = self.block_all(span, arena_vec![self; inner_match_stmt, yield_stmt], None);
diff --git a/src/librustc_ast_lowering/item.rs b/src/librustc_ast_lowering/item.rs
index 94ae2a0973a..c22b2812a9e 100644
--- a/src/librustc_ast_lowering/item.rs
+++ b/src/librustc_ast_lowering/item.rs
@@ -814,7 +814,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     }
 
     /// Construct `ExprKind::Err` for the given `span`.
-    fn expr_err(&mut self, span: Span) -> hir::Expr<'hir> {
+    crate fn expr_err(&mut self, span: Span) -> hir::Expr<'hir> {
         self.expr(span, hir::ExprKind::Err, AttrVec::new())
     }
 
@@ -960,7 +960,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         id
     }
 
-    fn lower_body(
+    pub(super) fn lower_body(
         &mut self,
         f: impl FnOnce(&mut Self) -> (&'hir [hir::Param<'hir>], hir::Expr<'hir>),
     ) -> hir::BodyId {
diff --git a/src/librustc_ast_lowering/lib.rs b/src/librustc_ast_lowering/lib.rs
index dcedcd51f50..24fe51fc4d0 100644
--- a/src/librustc_ast_lowering/lib.rs
+++ b/src/librustc_ast_lowering/lib.rs
@@ -116,6 +116,10 @@ struct LoweringContext<'a, 'hir: 'a> {
 
     generator_kind: Option<hir::GeneratorKind>,
 
+    /// When inside an `async` context, this is the `HirId` of the
+    /// `task_context` local bound to the resume argument of the generator.
+    task_context: Option<hir::HirId>,
+
     /// Used to get the current `fn`'s def span to point to when using `await`
     /// outside of an `async fn`.
     current_item: Option<Span>,
@@ -294,6 +298,7 @@ pub fn lower_crate<'a, 'hir>(
         item_local_id_counters: Default::default(),
         node_id_to_hir_id: IndexVec::new(),
         generator_kind: None,
+        task_context: None,
         current_item: None,
         lifetimes_to_define: Vec::new(),
         is_collecting_in_band_lifetimes: false,
diff --git a/src/librustc_mir/borrow_check/type_check/input_output.rs b/src/librustc_mir/borrow_check/type_check/input_output.rs
index 37cf77b7095..f2194c77c89 100644
--- a/src/librustc_mir/borrow_check/type_check/input_output.rs
+++ b/src/librustc_mir/borrow_check/type_check/input_output.rs
@@ -64,13 +64,16 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             }
         };
 
+        debug!(
+            "equate_inputs_and_outputs: normalized_input_tys = {:?}, local_decls = {:?}",
+            normalized_input_tys, body.local_decls
+        );
+
         // Equate expected input tys with those in the MIR.
         for (&normalized_input_ty, argument_index) in normalized_input_tys.iter().zip(0..) {
             // In MIR, argument N is stored in local N+1.
             let local = Local::new(argument_index + 1);
 
-            debug!("equate_inputs_and_outputs: normalized_input_ty = {:?}", normalized_input_ty);
-
             let mir_input_ty = body.local_decls[local].ty;
             let mir_input_span = body.local_decls[local].source_info.span;
             self.equate_normalized_input_or_output(
diff --git a/src/librustc_span/symbol.rs b/src/librustc_span/symbol.rs
index 5685505f694..35a97cfd3c3 100644
--- a/src/librustc_span/symbol.rs
+++ b/src/librustc_span/symbol.rs
@@ -544,7 +544,7 @@ symbols! {
         plugin_registrar,
         plugins,
         Poll,
-        poll_with_tls_context,
+        poll_with_context,
         powerpc_target_feature,
         precise_pointer_size_matching,
         pref_align_of,
@@ -720,6 +720,7 @@ symbols! {
         target_has_atomic_load_store,
         target_thread_local,
         task,
+        _task_context,
         tbm_target_feature,
         termination_trait,
         termination_trait_test,
diff --git a/src/libstd/future.rs b/src/libstd/future.rs
index c1ca6771326..c0675eeba98 100644
--- a/src/libstd/future.rs
+++ b/src/libstd/future.rs
@@ -1,12 +1,14 @@
 //! Asynchronous values.
 
-use core::cell::Cell;
-use core::marker::Unpin;
-use core::ops::{Drop, Generator, GeneratorState};
-use core::option::Option;
-use core::pin::Pin;
-use core::ptr::NonNull;
-use core::task::{Context, Poll};
+#[cfg(bootstrap)]
+use core::{
+    cell::Cell,
+    marker::Unpin,
+    ops::{Drop, Generator, GeneratorState},
+    pin::Pin,
+    ptr::NonNull,
+    task::{Context, Poll},
+};
 
 #[doc(inline)]
 #[stable(feature = "futures_api", since = "1.36.0")]
@@ -17,6 +19,7 @@ pub use core::future::*;
 /// This function returns a `GenFuture` underneath, but hides it in `impl Trait` to give
 /// better error messages (`impl Future` rather than `GenFuture<[closure.....]>`).
 // This is `const` to avoid extra errors after we recover from `const async fn`
+#[cfg(bootstrap)]
 #[doc(hidden)]
 #[unstable(feature = "gen_future", issue = "50547")]
 pub const fn from_generator<T: Generator<Yield = ()>>(x: T) -> impl Future<Output = T::Return> {
@@ -24,6 +27,7 @@ pub const fn from_generator<T: Generator<Yield = ()>>(x: T) -> impl Future<Outpu
 }
 
 /// A wrapper around generators used to implement `Future` for `async`/`await` code.
+#[cfg(bootstrap)]
 #[doc(hidden)]
 #[unstable(feature = "gen_future", issue = "50547")]
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
@@ -31,8 +35,10 @@ struct GenFuture<T: Generator<Yield = ()>>(T);
 
 // We rely on the fact that async/await futures are immovable in order to create
 // self-referential borrows in the underlying generator.
+#[cfg(bootstrap)]
 impl<T: Generator<Yield = ()>> !Unpin for GenFuture<T> {}
 
+#[cfg(bootstrap)]
 #[doc(hidden)]
 #[unstable(feature = "gen_future", issue = "50547")]
 impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
@@ -48,12 +54,15 @@ impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
     }
 }
 
+#[cfg(bootstrap)]
 thread_local! {
     static TLS_CX: Cell<Option<NonNull<Context<'static>>>> = Cell::new(None);
 }
 
+#[cfg(bootstrap)]
 struct SetOnDrop(Option<NonNull<Context<'static>>>);
 
+#[cfg(bootstrap)]
 impl Drop for SetOnDrop {
     fn drop(&mut self) {
         TLS_CX.with(|tls_cx| {
@@ -64,6 +73,7 @@ impl Drop for SetOnDrop {
 
 // Safety: the returned guard must drop before `cx` is dropped and before
 // any previous guard is dropped.
+#[cfg(bootstrap)]
 unsafe fn set_task_context(cx: &mut Context<'_>) -> SetOnDrop {
     // transmute the context's lifetime to 'static so we can store it.
     let cx = core::mem::transmute::<&mut Context<'_>, &mut Context<'static>>(cx);
@@ -71,6 +81,7 @@ unsafe fn set_task_context(cx: &mut Context<'_>) -> SetOnDrop {
     SetOnDrop(old_cx)
 }
 
+#[cfg(bootstrap)]
 #[doc(hidden)]
 #[unstable(feature = "gen_future", issue = "50547")]
 /// Polls a future in the current thread-local task waker.
diff --git a/src/test/ui/async-await/issues/issue-62009-1.stderr b/src/test/ui/async-await/issues/issue-62009-1.stderr
index cd6670923c2..0624c049048 100644
--- a/src/test/ui/async-await/issues/issue-62009-1.stderr
+++ b/src/test/ui/async-await/issues/issue-62009-1.stderr
@@ -33,10 +33,10 @@ error[E0277]: the trait bound `[closure@$DIR/issue-62009-1.rs:16:5: 16:15]: std:
 LL |     (|_| 2333).await;
    |     ^^^^^^^^^^^^^^^^ the trait `std::future::Future` is not implemented for `[closure@$DIR/issue-62009-1.rs:16:5: 16:15]`
    | 
-  ::: $SRC_DIR/libstd/future.rs:LL:COL
+  ::: $SRC_DIR/libcore/future/mod.rs:LL:COL
    |
 LL |     F: Future,
-   |        ------ required by this bound in `std::future::poll_with_tls_context`
+   |        ------ required by this bound in `std::future::poll_with_context`
 
 error: aborting due to 4 previous errors
 
diff --git a/src/test/ui/async-await/no-std.rs b/src/test/ui/async-await/no-std.rs
new file mode 100644
index 00000000000..63e93cdff7e
--- /dev/null
+++ b/src/test/ui/async-await/no-std.rs
@@ -0,0 +1,13 @@
+// edition:2018
+// check-pass
+
+#![no_std]
+#![crate_type = "rlib"]
+
+use core::future::Future;
+
+async fn a(f: impl Future) {
+    f.await;
+}
+
+fn main() {}