From dfcfa170f57695cd1169ae9dc65c0ed629e1a2a3 Mon Sep 17 00:00:00 2001 From: Jonas Schievink Date: Mon, 10 Feb 2020 17:23:09 +0100 Subject: [PATCH] Make async/await lowering use resume arguments --- src/librustc_ast_lowering/expr.rs | 95 +++++++++++++++---- src/librustc_ast_lowering/item.rs | 2 +- src/librustc_ast_lowering/lib.rs | 5 + src/librustc_span/symbol.rs | 3 +- .../async-await/issues/issue-62009-1.stderr | 4 +- 5 files changed, 86 insertions(+), 23 deletions(-) diff --git a/src/librustc_ast_lowering/expr.rs b/src/librustc_ast_lowering/expr.rs index a4cbae51966..104634f4fa0 100644 --- a/src/librustc_ast_lowering/expr.rs +++ b/src/librustc_ast_lowering/expr.rs @@ -10,6 +10,7 @@ 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]) -> &'hir [hir::Expr<'hir>] { @@ -483,14 +484,44 @@ pub(super) fn make_async_expr( 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 || -> { body }`: + // `static |task_context| -> { body }`: let generator_kind = hir::ExprKind::Closure( capture_clause, decl, @@ -523,9 +554,10 @@ pub(super) fn make_async_expr( /// ```rust /// match { /// 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 @@ fn lower_expr_await(&mut self, await_span: Span, expr: &Expr) -> hir::ExprKind<' 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 @@ fn lower_expr_await(&mut self, await_span: Span, expr: &Expr) -> hir::ExprKind<' "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 @@ fn lower_expr_await(&mut self, await_span: Span, expr: &Expr) -> hir::ExprKind<' 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), diff --git a/src/librustc_ast_lowering/item.rs b/src/librustc_ast_lowering/item.rs index aab6aa7c35b..b1fd3da99ce 100644 --- a/src/librustc_ast_lowering/item.rs +++ b/src/librustc_ast_lowering/item.rs @@ -809,7 +809,7 @@ fn lower_trait_item_ref(&mut self, i: &AssocItem) -> hir::TraitItemRef { } /// 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()) } diff --git a/src/librustc_ast_lowering/lib.rs b/src/librustc_ast_lowering/lib.rs index 32b0f0db358..3c1a82febda 100644 --- a/src/librustc_ast_lowering/lib.rs +++ b/src/librustc_ast_lowering/lib.rs @@ -117,6 +117,10 @@ struct LoweringContext<'a, 'hir: 'a> { generator_kind: Option, + /// 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, + /// Used to get the current `fn`'s def span to point to when using `await` /// outside of an `async fn`. current_item: Option, @@ -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, diff --git a/src/librustc_span/symbol.rs b/src/librustc_span/symbol.rs index 19754c83038..e7adbaa03eb 100644 --- a/src/librustc_span/symbol.rs +++ b/src/librustc_span/symbol.rs @@ -544,7 +544,7 @@ 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 @@ target_has_atomic_load_store, target_thread_local, task, + _task_context, tbm_target_feature, termination_trait, termination_trait_test, 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 -- 2.44.0