WIP: Remove ResumeTy from async lowering

Swatinem · Swatinem · commit e92f4f13a19a · 2023-09-24T15:27:44.000+02:00
Instead of using the stdlib supported `ResumeTy`, which is being converting to a `&mut Context<'_>` during the Generator MIR pass, this will use `&mut Context<'_>` directly in HIR lowering. It pretty much reverts rust-lang#105977 and re-applies an updated version of rust-lang#105250. This still fails the testcase added in rust-lang#106264 however, for reasons I don’t understand.
diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
@@ -16,7 +16,7 @@ use rustc_hir::def::Res;
 use rustc_hir::definitions::DefPathData;
 use rustc_session::errors::report_lit_error;
 use rustc_span::source_map::{respan, DesugaringKind, Span, Spanned};
-use rustc_span::symbol::{sym, Ident, Symbol};
+use rustc_span::symbol::{kw, sym, Ident, Symbol};
 use rustc_span::DUMMY_SP;
 use thin_vec::{thin_vec, ThinVec};
 
@@ -603,14 +603,28 @@ impl<'hir> LoweringContext<'_, 'hir> {
     ) -> hir::ExprKind<'hir> {
         let output = ret_ty.unwrap_or_else(|| hir::FnRetTy::DefaultReturn(self.lower_span(span)));
 
-        // Resume argument type: `ResumeTy`
-        let unstable_span =
-            self.mark_span_with_reason(DesugaringKind::Async, span, self.allow_gen_future.clone());
-        let resume_ty = hir::QPath::LangItem(hir::LangItem::ResumeTy, unstable_span, None);
+        // Resume argument type: `&mut Context<'_>`.
+        let context_lifetime_ident = Ident::with_dummy_span(kw::UnderscoreLifetime);
+        let context_lifetime = self.arena.alloc(hir::Lifetime {
+            hir_id: self.next_id(),
+            ident: context_lifetime_ident,
+            res: hir::LifetimeName::Infer,
+        });
+        let context_path =
+            hir::QPath::LangItem(hir::LangItem::Context, self.lower_span(span), None);
+        let context_ty = hir::MutTy {
+            ty: self.arena.alloc(hir::Ty {
+                hir_id: self.next_id(),
+                kind: hir::TyKind::Path(context_path),
+                span: self.lower_span(span),
+            }),
+            mutbl: hir::Mutability::Mut,
+        };
+
         let input_ty = hir::Ty {
             hir_id: self.next_id(),
-            kind: hir::TyKind::Path(resume_ty),
-            span: unstable_span,
+            kind: hir::TyKind::Ref(context_lifetime, context_ty),
+            span: self.lower_span(span),
         };
 
         // The closure/generator `FnDecl` takes a single (resume) argument of type `input_ty`.
@@ -693,7 +707,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     ///     mut __awaitee => loop {
     ///         match unsafe { ::std::future::Future::poll(
     ///             <::std::pin::Pin>::new_unchecked(&mut __awaitee),
-    ///             ::std::future::get_context(task_context),
+    ///             task_context,
     ///         ) } {
     ///             ::std::task::Poll::Ready(result) => break result,
     ///             ::std::task::Poll::Pending => {}
@@ -714,11 +728,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
             }
         }
         let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, None);
-        let gen_future_span = self.mark_span_with_reason(
-            DesugaringKind::Await,
-            full_span,
-            self.allow_gen_future.clone(),
-        );
         let expr = self.lower_expr_mut(expr);
         let expr_hir_id = expr.hir_id;
 
@@ -734,7 +743,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         // unsafe {
         //     ::std::future::Future::poll(
         //         ::std::pin::Pin::new_unchecked(&mut __awaitee),
-        //         ::std::future::get_context(task_context),
+        //         task_context,
         //     )
         // }
         let poll_expr = {
@@ -752,16 +761,10 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 arena_vec![self; ref_mut_awaitee],
                 Some(expr_hir_id),
             );
-            let get_context = self.expr_call_lang_item_fn_mut(
-                gen_future_span,
-                hir::LangItem::GetContext,
-                arena_vec![self; task_context],
-                Some(expr_hir_id),
-            );
             let call = self.expr_call_lang_item_fn(
                 span,
                 hir::LangItem::FuturePoll,
-                arena_vec![self; new_unchecked, get_context],
+                arena_vec![self; new_unchecked, task_context],
                 Some(expr_hir_id),
             );
             self.arena.alloc(self.expr_unsafe(call))
@@ -772,9 +775,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let loop_hir_id = self.lower_node_id(loop_node_id);
         let ready_arm = {
             let x_ident = Ident::with_dummy_span(sym::result);
-            let (x_pat, x_pat_hid) = self.pat_ident(gen_future_span, x_ident);
-            let x_expr = self.expr_ident(gen_future_span, x_ident, x_pat_hid);
-            let ready_field = self.single_pat_field(gen_future_span, x_pat);
+            let (x_pat, x_pat_hid) = self.pat_ident(full_span, x_ident);
+            let x_expr = self.expr_ident(full_span, x_ident, x_pat_hid);
+            let ready_field = self.single_pat_field(full_span, x_pat);
             let ready_pat = self.pat_lang_item_variant(
                 span,
                 hir::LangItem::PollReady,
@@ -784,7 +787,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             let break_x = self.with_loop_scope(loop_node_id, move |this| {
                 let expr_break =
                     hir::ExprKind::Break(this.lower_loop_destination(None), Some(x_expr));
-                this.arena.alloc(this.expr(gen_future_span, expr_break))
+                this.arena.alloc(this.expr(full_span, expr_break))
             });
             self.arm(ready_pat, break_x)
         };
diff --git a/compiler/rustc_hir/src/lang_items.rs b/compiler/rustc_hir/src/lang_items.rs
@@ -294,11 +294,6 @@ language_item_table! {
     PollReady,               sym::Ready,               poll_ready_variant,         Target::Variant,        GenericRequirement::None;
     PollPending,             sym::Pending,             poll_pending_variant,       Target::Variant,        GenericRequirement::None;
 
-    // FIXME(swatinem): the following lang items are used for async lowering and
-    // should become obsolete eventually.
-    ResumeTy,                sym::ResumeTy,            resume_ty,                  Target::Struct,         GenericRequirement::None;
-    GetContext,              sym::get_context,         get_context_fn,             Target::Fn,             GenericRequirement::None;
-
     Context,                 sym::Context,             context,                    Target::Struct,         GenericRequirement::None;
     FuturePoll,              sym::poll,                future_poll_fn,             Target::Method(MethodKind::Trait { body: false }), GenericRequirement::None;
 
diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
@@ -465,104 +465,6 @@ fn replace_local<'tcx>(
     new_local
 }
 
-/// Transforms the `body` of the generator applying the following transforms:
-///
-/// - Eliminates all the `get_context` calls that async lowering created.
-/// - Replace all `Local` `ResumeTy` types with `&mut Context<'_>` (`context_mut_ref`).
-///
-/// The `Local`s that have their types replaced are:
-/// - The `resume` argument itself.
-/// - The argument to `get_context`.
-/// - The yielded value of a `yield`.
-///
-/// The `ResumeTy` hides a `&mut Context<'_>` behind an unsafe raw pointer, and the
-/// `get_context` function is being used to convert that back to a `&mut Context<'_>`.
-///
-/// Ideally the async lowering would not use the `ResumeTy`/`get_context` indirection,
-/// but rather directly use `&mut Context<'_>`, however that would currently
-/// lead to higher-kinded lifetime errors.
-/// See <https://github.com/rust-lang/rust/issues/105501>.
-///
-/// The async lowering step and the type / lifetime inference / checking are
-/// still using the `ResumeTy` indirection for the time being, and that indirection
-/// is removed here. After this transform, the generator body only knows about `&mut Context<'_>`.
-fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    let context_mut_ref = Ty::new_task_context(tcx);
-
-    // replace the type of the `resume` argument
-    replace_resume_ty_local(tcx, body, Local::new(2), context_mut_ref);
-
-    let get_context_def_id = tcx.require_lang_item(LangItem::GetContext, None);
-
-    for bb in START_BLOCK..body.basic_blocks.next_index() {
-        let bb_data = &body[bb];
-        if bb_data.is_cleanup {
-            continue;
-        }
-
-        match &bb_data.terminator().kind {
-            TerminatorKind::Call { func, .. } => {
-                let func_ty = func.ty(body, tcx);
-                if let ty::FnDef(def_id, _) = *func_ty.kind() {
-                    if def_id == get_context_def_id {
-                        let local = eliminate_get_context_call(&mut body[bb]);
-                        replace_resume_ty_local(tcx, body, local, context_mut_ref);
-                    }
-                } else {
-                    continue;
-                }
-            }
-            TerminatorKind::Yield { resume_arg, .. } => {
-                replace_resume_ty_local(tcx, body, resume_arg.local, context_mut_ref);
-            }
-            _ => {}
-        }
-    }
-}
-
-fn eliminate_get_context_call<'tcx>(bb_data: &mut BasicBlockData<'tcx>) -> Local {
-    let terminator = bb_data.terminator.take().unwrap();
-    if let TerminatorKind::Call { mut args, destination, target, .. } = terminator.kind {
-        let arg = args.pop().unwrap();
-        let local = arg.place().unwrap().local;
-
-        let arg = Rvalue::Use(arg);
-        let assign = Statement {
-            source_info: terminator.source_info,
-            kind: StatementKind::Assign(Box::new((destination, arg))),
-        };
-        bb_data.statements.push(assign);
-        bb_data.terminator = Some(Terminator {
-            source_info: terminator.source_info,
-            kind: TerminatorKind::Goto { target: target.unwrap() },
-        });
-        local
-    } else {
-        bug!();
-    }
-}
-
-#[cfg_attr(not(debug_assertions), allow(unused))]
-fn replace_resume_ty_local<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    body: &mut Body<'tcx>,
-    local: Local,
-    context_mut_ref: Ty<'tcx>,
-) {
-    let local_ty = std::mem::replace(&mut body.local_decls[local].ty, context_mut_ref);
-    // We have to replace the `ResumeTy` that is used for type and borrow checking
-    // with `&mut Context<'_>` in MIR.
-    #[cfg(debug_assertions)]
-    {
-        if let ty::Adt(resume_ty_adt, _) = local_ty.kind() {
-            let expected_adt = tcx.adt_def(tcx.require_lang_item(LangItem::ResumeTy, None));
-            assert_eq!(*resume_ty_adt, expected_adt);
-        } else {
-            panic!("expected `ResumeTy`, found `{:?}`", local_ty);
-        };
-    }
-}
-
 struct LivenessInfo {
     /// Which locals are live across any suspension point.
     saved_locals: GeneratorSavedLocals,
@@ -1458,22 +1360,13 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         // RETURN_PLACE then is a fresh unused local with type ret_ty.
         let new_ret_local = replace_local(RETURN_PLACE, ret_ty, body, tcx);
 
-        // Replace all occurrences of `ResumeTy` with `&mut Context<'_>` within async bodies.
-        if is_async_kind {
-            transform_async_context(tcx, body);
-        }
-
         // We also replace the resume argument and insert an `Assign`.
         // This is needed because the resume argument `_2` might be live across a `yield`, in which
         // case there is no `Assign` to it that the transform can turn into a store to the generator
         // state. After the yield the slot in the generator state would then be uninitialized.
         let resume_local = Local::new(2);
-        let resume_ty = if is_async_kind {
-            Ty::new_task_context(tcx)
-        } else {
-            body.local_decls[resume_local].ty
-        };
-        let new_resume_local = replace_local(resume_local, resume_ty, body, tcx);
+        let new_resume_local =
+            replace_local(resume_local, body.local_decls[resume_local].ty, body, tcx);
 
         // When first entering the generator, move the resume argument into its new local.
         let source_info = SourceInfo::outermost(body.span);
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
@@ -271,7 +271,6 @@ symbols! {
         Relaxed,
         Release,
         Result,
-        ResumeTy,
         Return,
         Right,
         Rust,
@@ -801,7 +800,6 @@ symbols! {
         generic_const_exprs,
         generic_const_items,
         generic_param_attrs,
-        get_context,
         global_allocator,
         global_asm,
         globs,
diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs
@@ -120,41 +120,23 @@ fn fn_sig_for_fn_abi<'tcx>(
             // `Generator::resume(...) -> GeneratorState` function in case we
             // have an ordinary generator, or the `Future::poll(...) -> Poll`
             // function in case this is a special generator backing an async construct.
-            let (resume_ty, ret_ty) = if tcx.generator_is_async(did) {
-                // The signature should be `Future::poll(_, &mut Context<'_>) -> Poll<Output>`
+            let ret_ty = if tcx.generator_is_async(did) {
+                // The signature should be `Future::poll(..) -> Poll<Output>`
                 let poll_did = tcx.require_lang_item(LangItem::Poll, None);
                 let poll_adt_ref = tcx.adt_def(poll_did);
                 let poll_args = tcx.mk_args(&[sig.return_ty.into()]);
-                let ret_ty = Ty::new_adt(tcx, poll_adt_ref, poll_args);
-
-                // We have to replace the `ResumeTy` that is used for type and borrow checking
-                // with `&mut Context<'_>` which is used in codegen.
-                #[cfg(debug_assertions)]
-                {
-                    if let ty::Adt(resume_ty_adt, _) = sig.resume_ty.kind() {
-                        let expected_adt =
-                            tcx.adt_def(tcx.require_lang_item(LangItem::ResumeTy, None));
-                        assert_eq!(*resume_ty_adt, expected_adt);
-                    } else {
-                        panic!("expected `ResumeTy`, found `{:?}`", sig.resume_ty);
-                    };
-                }
-                let context_mut_ref = Ty::new_task_context(tcx);
-
-                (context_mut_ref, ret_ty)
+                Ty::new_adt(tcx, poll_adt_ref, poll_args)
             } else {
-                // The signature should be `Generator::resume(_, Resume) -> GeneratorState<Yield, Return>`
+                // The signature should be `Generator::resume(..) -> GeneratorState<Yield, Return>`
                 let state_did = tcx.require_lang_item(LangItem::GeneratorState, None);
                 let state_adt_ref = tcx.adt_def(state_did);
                 let state_args = tcx.mk_args(&[sig.yield_ty.into(), sig.return_ty.into()]);
-                let ret_ty = Ty::new_adt(tcx, state_adt_ref, state_args);
-
-                (sig.resume_ty, ret_ty)
+                Ty::new_adt(tcx, state_adt_ref, state_args)
             };
 
             ty::Binder::bind_with_vars(
                 tcx.mk_fn_sig(
-                    [env_ty, resume_ty],
+                    [env_ty, sig.resume_ty],
                     ret_ty,
                     false,
                     hir::Unsafety::Normal,
diff --git a/library/core/src/future/mod.rs b/library/core/src/future/mod.rs
@@ -44,7 +44,7 @@ pub use poll_fn::{poll_fn, PollFn};
 ///    non-Send/Sync as well, and we don't want that.
 ///
 /// It also simplifies the HIR lowering of `.await`.
-#[lang = "ResumeTy"]
+#[cfg_attr(bootstrap, lang = "ResumeTy")]
 #[doc(hidden)]
 #[unstable(feature = "gen_future", issue = "50547")]
 #[derive(Debug, Copy, Clone)]
@@ -56,7 +56,7 @@ unsafe impl Send for ResumeTy {}
 #[unstable(feature = "gen_future", issue = "50547")]
 unsafe impl Sync for ResumeTy {}
 
-#[lang = "get_context"]
+#[cfg_attr(bootstrap, lang = "get_context")]
 #[doc(hidden)]
 #[unstable(feature = "gen_future", issue = "50547")]
 #[must_use]
