Auto merge of #70831 - sfackler:shrink-future-stack, r=matthewjasper

Remove a stack frame from .await calls

The stack frames when `.await`ing one async fn from another currently look like this:

```
  12: foo::b::{{closure}}
             at src/main.rs:2
  13: <core::future::from_generator::GenFuture<T> as core::future::future::Future>::poll
             at /home/sfackler/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/src/libcore/future/mod.rs:66
  14: core::future::poll_with_context
             at /home/sfackler/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/src/libcore/future/mod.rs:84
  15: foo::a::{{closure}}
             at src/main.rs:6
```

Since the move away from using TLS to pass the Context around, it's now easy to remove frame 14 by removing poll_with_context in favor of calling Future::poll directly. This still leaves the `GenFuture` frame, but that seems significantly harder to deal with.

It also improves diagnostics a bit since they no longer talk about the private poll_with_context function.

5 files changed, 40 insertions, 25 deletions

diff --git a/src/libcore/future/mod.rs b/src/libcore/future/mod.rs
index 8dfda7a4a32..a6b769147d0 100644
--- a/src/libcore/future/mod.rs
+++ b/src/libcore/future/mod.rs
@@ -77,9 +77,6 @@ where
 #[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())
+pub unsafe fn get_context<'a, 'b>(cx: ResumeTy) -> &'a mut Context<'b> {
+    &mut *cx.0.as_ptr().cast()
 }
diff --git a/src/librustc_ast_lowering/expr.rs b/src/librustc_ast_lowering/expr.rs
index 2fc58efbf0e..0eed47050aa 100644
--- a/src/librustc_ast_lowering/expr.rs
+++ b/src/librustc_ast_lowering/expr.rs
@@ -556,9 +556,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
     /// ```rust
     /// match <expr> {
     ///     mut pinned => loop {
-    ///         match unsafe { ::std::future::poll_with_context(
+    ///         match unsafe { ::std::future::Future::poll(
     ///             <::std::pin::Pin>::new_unchecked(&mut pinned),
-    ///             task_context,
+    ///             ::std::future::get_context(task_context),
     ///         ) } {
     ///             ::std::task::Poll::Ready(result) => break result,
     ///             ::std::task::Poll::Pending => {}
@@ -599,9 +599,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let task_context_ident = Ident::with_dummy_span(sym::_task_context);
 
         // unsafe {
-        //     ::std::future::poll_with_context(
+        //     ::std::future::Future::poll(
         //         ::std::pin::Pin::new_unchecked(&mut pinned),
-        //         task_context,
+        //         ::std::future::get_context(task_context),
         //     )
         // }
         let poll_expr = {
@@ -622,10 +622,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 arena_vec![self; ref_mut_pinned],
             );
             let new_unchecked = self.expr(span, new_unchecked_expr_kind, ThinVec::new());
-            let call = self.expr_call_std_path(
+            let get_context = self.expr_call_std_path_mut(
                 gen_future_span,
-                &[sym::future, sym::poll_with_context],
-                arena_vec![self; new_unchecked, task_context],
+                &[sym::future, sym::get_context],
+                arena_vec![self; task_context],
+            );
+            let call = self.expr_call_std_path(
+                span,
+                &[sym::future, sym::Future, sym::poll],
+                arena_vec![self; new_unchecked, get_context],
             );
             self.arena.alloc(self.expr_unsafe(call))
         };
@@ -1326,25 +1331,43 @@ impl<'hir> LoweringContext<'_, 'hir> {
         self.arena.alloc(self.expr(sp, hir::ExprKind::Tup(&[]), ThinVec::new()))
     }
 
+    fn expr_call_mut(
+        &mut self,
+        span: Span,
+        e: &'hir hir::Expr<'hir>,
+        args: &'hir [hir::Expr<'hir>],
+    ) -> hir::Expr<'hir> {
+        self.expr(span, hir::ExprKind::Call(e, args), ThinVec::new())
+    }
+
     fn expr_call(
         &mut self,
         span: Span,
         e: &'hir hir::Expr<'hir>,
         args: &'hir [hir::Expr<'hir>],
     ) -> &'hir hir::Expr<'hir> {
-        self.arena.alloc(self.expr(span, hir::ExprKind::Call(e, args), ThinVec::new()))
+        self.arena.alloc(self.expr_call_mut(span, e, args))
     }
 
     // Note: associated functions must use `expr_call_std_path`.
-    fn expr_call_std_path(
+    fn expr_call_std_path_mut(
         &mut self,
         span: Span,
         path_components: &[Symbol],
         args: &'hir [hir::Expr<'hir>],
-    ) -> &'hir hir::Expr<'hir> {
+    ) -> hir::Expr<'hir> {
         let path =
             self.arena.alloc(self.expr_std_path(span, path_components, None, ThinVec::new()));
-        self.expr_call(span, path, args)
+        self.expr_call_mut(span, path, args)
+    }
+
+    fn expr_call_std_path(
+        &mut self,
+        span: Span,
+        path_components: &[Symbol],
+        args: &'hir [hir::Expr<'hir>],
+    ) -> &'hir hir::Expr<'hir> {
+        self.arena.alloc(self.expr_call_std_path_mut(span, path_components, args))
     }
 
     // Create an expression calling an associated function of an std type.
diff --git a/src/librustc_span/symbol.rs b/src/librustc_span/symbol.rs
index 4c963ac84dc..1575e6fd533 100644
--- a/src/librustc_span/symbol.rs
+++ b/src/librustc_span/symbol.rs
@@ -348,6 +348,7 @@ symbols! {
         generators,
         generic_associated_types,
         generic_param_attrs,
+        get_context,
         global_allocator,
         global_asm,
         globs,
@@ -547,8 +548,8 @@ symbols! {
         plugin,
         plugin_registrar,
         plugins,
+        poll,
         Poll,
-        poll_with_context,
         powerpc_target_feature,
         precise_pointer_size_matching,
         pref_align_of,
diff --git a/src/test/ui/async-await/issue-70594.stderr b/src/test/ui/async-await/issue-70594.stderr
index d2fa7e58f6a..496ca506c60 100644
--- a/src/test/ui/async-await/issue-70594.stderr
+++ b/src/test/ui/async-await/issue-70594.stderr
@@ -32,11 +32,8 @@ error[E0277]: the trait bound `(): std::future::Future` is not satisfied
    |
 LL |     [1; ().await];
    |         ^^^^^^^^ the trait `std::future::Future` is not implemented for `()`
-   | 
-  ::: $SRC_DIR/libcore/future/mod.rs:LL:COL
    |
-LL |     F: Future,
-   |        ------ required by this bound in `std::future::poll_with_context`
+   = note: required by `std::future::Future::poll`
 
 error: aborting due to 5 previous errors
 
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 2417b592c7d..ec4e9e397a8 100644
--- a/src/test/ui/async-await/issues/issue-62009-1.stderr
+++ b/src/test/ui/async-await/issues/issue-62009-1.stderr
@@ -32,11 +32,8 @@ error[E0277]: the trait bound `[closure@$DIR/issue-62009-1.rs:12:5: 12:15]: std:
    |
 LL |     (|_| 2333).await;
    |     ^^^^^^^^^^^^^^^^ the trait `std::future::Future` is not implemented for `[closure@$DIR/issue-62009-1.rs:12:5: 12:15]`
-   | 
-  ::: $SRC_DIR/libcore/future/mod.rs:LL:COL
    |
-LL |     F: Future,
-   |        ------ required by this bound in `std::future::poll_with_context`
+   = note: required by `std::future::Future::poll`
 
 error: aborting due to 4 previous errors