Make async/await lowering use resume arguments

Author: jonas-schievink

src/librustc_ast_lowering/expr.rs

@@ -10,6 +10,7 @@ use rustc_hir as hir;
 use rustc_hir::def::Res;
 use rustc_span::source_map::{respan, DesugaringKind, Span, Spanned};
 use rustc_span::symbol::{sym, Symbol};
+use rustc_span::DUMMY_SP;
 
 impl<'hir> LoweringContext<'_, 'hir> {
     fn lower_exprs(&mut self, exprs: &[AstP<Expr>]) -> &'hir [hir::Expr<'hir>] {
@@ -483,14 +484,44 @@ impl<'hir> LoweringContext<'_, 'hir> {
             Some(ty) => FnRetTy::Ty(ty),
             None => FnRetTy::Default(span),
         };
-        let ast_decl = FnDecl { inputs: vec![], output };
+
+        let task_context_id = self.resolver.next_node_id();
+        let task_context_hid = self.lower_node_id(task_context_id);
+        let ast_decl = FnDecl {
+            inputs: vec![Param {
+                attrs: AttrVec::new(),
+                ty: AstP(Ty {
+                    id: self.resolver.next_node_id(),
+                    kind: TyKind::Infer,
+                    span: DUMMY_SP,
+                }),
+                pat: AstP(Pat {
+                    id: task_context_id,
+                    kind: PatKind::Ident(
+                        BindingMode::ByValue(Mutability::Mut),
+                        Ident::with_dummy_span(sym::_task_context),
+                        None,
+                    ),
+                    span: DUMMY_SP,
+                }),
+                id: self.resolver.next_node_id(),
+                span: DUMMY_SP,
+                is_placeholder: false,
+            }],
+            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| {
             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;
+            res
         });
 
-        // `static || -> <ret_ty> { body }`:
+        // `static |task_context| -> <ret_ty> { body }`:
         let generator_kind = hir::ExprKind::Closure(
             capture_clause,
             decl,
@@ -523,9 +554,10 @@ 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 => {}
     ///         }
@@ -561,12 +593,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 +618,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`
@@ -629,12 +671,27 @@ 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);
+        let loop_block =
+            self.block_all(span, arena_vec![self; inner_match_stmt, yield_stmt], None);
 
-        // loop { .. }
+        // loop { ...; task_context = yield (); }
         let loop_expr = self.arena.alloc(hir::Expr {
             hir_id: loop_hir_id,
             kind: hir::ExprKind::Loop(loop_block, None, hir::LoopSource::Loop),

src/librustc_ast_lowering/item.rs

@@ -809,7 +809,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())
     }
 

src/librustc_ast_lowering/lib.rs

@@ -117,6 +117,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>,
@@ -295,6 +299,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,

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,

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