Rollup merge of #105977 - Swatinem:async-mir-context, r=oli-obk

Transform async `ResumeTy` in generator transform

- Eliminates all the `get_context` calls that async lowering created.
- Replace all `Local` `ResumeTy` types with `&mut Context<'_>`.

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<'_>`.

---

Fixes https://github.com/bjorn3/rustc_codegen_cranelift/issues/1330 CC `@bjorn3`

r? `@compiler-errors`

compiler/rustc_hir/src/lang_items.rs

@@ -291,6 +291,7 @@ language_item_table! {
     IdentityFuture,          sym::identity_future,     identity_future_fn,         Target::Fn,             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;
 
     FromFrom,                sym::from,                from_fn,                    Target::Method(MethodKind::Trait { body: false }), GenericRequirement::None;

compiler/rustc_middle/src/ty/context.rs

@@ -1952,6 +1952,15 @@ impl<'tcx> TyCtxt<'tcx> {
         self.mk_ty(GeneratorWitness(types))
     }
 
+    /// Creates a `&mut Context<'_>` [`Ty`] with erased lifetimes.
+    pub fn mk_task_context(self) -> Ty<'tcx> {
+        let context_did = self.require_lang_item(LangItem::Context, None);
+        let context_adt_ref = self.adt_def(context_did);
+        let context_substs = self.intern_substs(&[self.lifetimes.re_erased.into()]);
+        let context_ty = self.mk_adt(context_adt_ref, context_substs);
+        self.mk_mut_ref(self.lifetimes.re_erased, context_ty)
+    }
+
     #[inline]
     pub fn mk_ty_var(self, v: TyVid) -> Ty<'tcx> {
         self.mk_ty_infer(TyVar(v))

compiler/rustc_mir_transform/src/generator.rs

@@ -460,6 +460,104 @@ 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 = tcx.mk_task_context();
+
+    // 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 BasicBlock::new(0)..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,
@@ -1283,13 +1381,13 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
             }
         };
 
-        let is_async_kind = body.generator_kind().unwrap() != GeneratorKind::Gen;
+        let is_async_kind = matches!(body.generator_kind(), Some(GeneratorKind::Async(_)));
         let (state_adt_ref, state_substs) = if is_async_kind {
             // Compute Poll<return_ty>
-            let state_did = tcx.require_lang_item(LangItem::Poll, None);
+            let poll_did = tcx.require_lang_item(LangItem::Poll, None);
-            let state_adt_ref = tcx.adt_def(state_did);
+            let poll_adt_ref = tcx.adt_def(poll_did);
-            let state_substs = tcx.intern_substs(&[body.return_ty().into()]);
+            let poll_substs = tcx.intern_substs(&[body.return_ty().into()]);
-            (state_adt_ref, state_substs)
+            (poll_adt_ref, poll_substs)
         } else {
             // Compute GeneratorState<yield_ty, return_ty>
             let state_did = tcx.require_lang_item(LangItem::GeneratorState, None);
@@ -1303,13 +1401,19 @@ 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 new_resume_local =
+        let resume_ty =
-            replace_local(resume_local, body.local_decls[resume_local].ty, body, tcx);
+            if is_async_kind { tcx.mk_task_context() } else { body.local_decls[resume_local].ty };
+        let new_resume_local = replace_local(resume_local, resume_ty, body, tcx);
 
         // When first entering the generator, move the resume argument into its new local.
         let source_info = SourceInfo::outermost(body.span);

compiler/rustc_span/src/symbol.rs

@@ -164,6 +164,7 @@ symbols! {
         Capture,
         Center,
         Clone,
+        Context,
         Continue,
         Copy,
         Count,

compiler/rustc_ty_utils/src/abi.rs

@@ -108,21 +108,41 @@ 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 ret_ty = if tcx.generator_is_async(did) {
+            let (resume_ty, ret_ty) = if tcx.generator_is_async(did) {
-                let state_did = tcx.require_lang_item(LangItem::Poll, None);
+                // The signature should be `Future::poll(_, &mut Context<'_>) -> Poll<Output>`
-                let state_adt_ref = tcx.adt_def(state_did);
+                let poll_did = tcx.require_lang_item(LangItem::Poll, None);
-                let state_substs = tcx.intern_substs(&[sig.return_ty.into()]);
+                let poll_adt_ref = tcx.adt_def(poll_did);
-                tcx.mk_adt(state_adt_ref, state_substs)
+                let poll_substs = tcx.intern_substs(&[sig.return_ty.into()]);
+                let ret_ty = tcx.mk_adt(poll_adt_ref, poll_substs);
+
+                // 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 = tcx.mk_task_context();
+
+                (context_mut_ref, ret_ty)
             } else {
+                // The signature should be `Generator::resume(_, 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_substs = tcx.intern_substs(&[sig.yield_ty.into(), sig.return_ty.into()]);
-                tcx.mk_adt(state_adt_ref, state_substs)
+                let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
+
+                (sig.resume_ty, ret_ty)
             };
 
             ty::Binder::bind_with_vars(
                 tcx.mk_fn_sig(
-                    [env_ty, sig.resume_ty].iter(),
+                    [env_ty, resume_ty].iter(),
                     &ret_ty,
                     false,
                     hir::Unsafety::Normal,

library/core/src/future/mod.rs

@@ -112,6 +112,10 @@ pub unsafe fn get_context<'a, 'b>(cx: ResumeTy) -> &'a mut Context<'b> {
     unsafe { &mut *cx.0.as_ptr().cast() }
 }
 
+// FIXME(swatinem): This fn is currently needed to work around shortcomings
+// in type and lifetime inference.
+// See the comment at the bottom of `LoweringContext::make_async_expr` and
+// <https://github.com/rust-lang/rust/issues/104826>.
 #[doc(hidden)]
 #[unstable(feature = "gen_future", issue = "50547")]
 #[inline]

library/core/src/task/wake.rs

@@ -174,6 +174,7 @@ impl RawWakerVTable {
 /// Currently, `Context` only serves to provide access to a [`&Waker`](Waker)
 /// which can be used to wake the current task.
 #[stable(feature = "futures_api", since = "1.36.0")]
+#[cfg_attr(not(bootstrap), lang = "Context")]
 pub struct Context<'a> {
     waker: &'a Waker,
     // Ensure we future-proof against variance changes by forcing

tests/mir-opt/building/async_await.a-{closure#0}.generator_resume.0.mir

@@ -0,0 +1,41 @@
+// MIR for `a::{closure#0}` 0 generator_resume
+/* generator_layout = GeneratorLayout {
+    field_tys: {},
+    variant_fields: {
+        Unresumed(0): [],
+        Returned (1): [],
+        Panicked (2): [],
+    },
+    storage_conflicts: BitMatrix(0x0) {},
+} */
+
+fn a::{closure#0}(_1: Pin<&mut [async fn body@$DIR/async_await.rs:11:14: 11:16]>, _2: &mut Context<'_>) -> Poll<()> {
+    debug _task_context => _4;           // in scope 0 at $DIR/async_await.rs:+0:14: +0:16
+    let mut _0: std::task::Poll<()>;     // return place in scope 0 at $DIR/async_await.rs:+0:14: +0:16
+    let mut _3: ();                      // in scope 0 at $DIR/async_await.rs:+0:14: +0:16
+    let mut _4: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+0:14: +0:16
+    let mut _5: u32;                     // in scope 0 at $DIR/async_await.rs:+0:14: +0:16
+
+    bb0: {
+        _5 = discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:11:14: 11:16]))); // scope 0 at $DIR/async_await.rs:+0:14: +0:16
+        switchInt(move _5) -> [0: bb1, 1: bb2, otherwise: bb3]; // scope 0 at $DIR/async_await.rs:+0:14: +0:16
+    }
+
+    bb1: {
+        _4 = move _2;                    // scope 0 at $DIR/async_await.rs:+0:14: +0:16
+        _3 = const ();                   // scope 0 at $DIR/async_await.rs:+0:14: +0:16
+        Deinit(_0);                      // scope 0 at $DIR/async_await.rs:+0:16: +0:16
+        ((_0 as Ready).0: ()) = move _3; // scope 0 at $DIR/async_await.rs:+0:16: +0:16
+        discriminant(_0) = 0;            // scope 0 at $DIR/async_await.rs:+0:16: +0:16
+        discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:11:14: 11:16]))) = 1; // scope 0 at $DIR/async_await.rs:+0:16: +0:16
+        return;                          // scope 0 at $DIR/async_await.rs:+0:16: +0:16
+    }
+
+    bb2: {
+        assert(const false, "`async fn` resumed after completion") -> bb2; // scope 0 at $DIR/async_await.rs:+0:14: +0:16
+    }
+
+    bb3: {
+        unreachable;                     // scope 0 at $DIR/async_await.rs:+0:14: +0:16
+    }
+}

tests/mir-opt/building/async_await.b-{closure#0}.generator_resume.0.mir

@@ -0,0 +1,337 @@
+// MIR for `b::{closure#0}` 0 generator_resume
+/* generator_layout = GeneratorLayout {
+    field_tys: {
+        _0: impl std::future::Future<Output = ()>,
+        _1: impl std::future::Future<Output = ()>,
+    },
+    variant_fields: {
+        Unresumed(0): [],
+        Returned (1): [],
+        Panicked (2): [],
+        Suspend0 (3): [_0],
+        Suspend1 (4): [_1],
+    },
+    storage_conflicts: BitMatrix(2x2) {
+        (_0, _0),
+        (_1, _1),
+    },
+} */
+
+fn b::{closure#0}(_1: Pin<&mut [async fn body@$DIR/async_await.rs:14:18: 17:2]>, _2: &mut Context<'_>) -> Poll<()> {
+    debug _task_context => _38;          // in scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    let mut _0: std::task::Poll<()>;     // return place in scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    let _3: ();                          // in scope 0 at $DIR/async_await.rs:+1:5: +1:14
+    let mut _4: impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _5: impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+1:5: +1:8
+    let mut _6: impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _7: ();                      // in scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    let _8: ();                          // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _9: std::task::Poll<()>;     // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _10: std::pin::Pin<&mut impl std::future::Future<Output = ()>>; // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _11: &mut impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _12: &mut impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _13: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+1:5: +1:14
+    let mut _14: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+1:5: +1:14
+    let mut _15: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _16: isize;                  // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _18: !;                      // in scope 0 at $DIR/async_await.rs:+1:5: +1:14
+    let mut _19: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _20: ();                     // in scope 0 at $DIR/async_await.rs:+1:8: +1:14
+    let mut _21: impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _22: impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+2:5: +2:8
+    let mut _23: impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let _24: ();                         // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _25: std::task::Poll<()>;    // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _26: std::pin::Pin<&mut impl std::future::Future<Output = ()>>; // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _27: &mut impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _28: &mut impl std::future::Future<Output = ()>; // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _29: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+2:5: +2:14
+    let mut _30: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+2:5: +2:14
+    let mut _31: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _32: isize;                  // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _34: !;                      // in scope 0 at $DIR/async_await.rs:+2:5: +2:14
+    let mut _35: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _36: ();                     // in scope 0 at $DIR/async_await.rs:+2:8: +2:14
+    let mut _37: ();                     // in scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    let mut _38: &mut std::task::Context<'_>; // in scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    let mut _39: u32;                    // in scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    scope 1 {
+        debug __awaitee => (((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2])) as variant#3).0: impl std::future::Future<Output = ()>); // in scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        let _17: ();                     // in scope 1 at $DIR/async_await.rs:+1:5: +1:14
+        scope 2 {
+        }
+        scope 3 {
+            debug result => _17;         // in scope 3 at $DIR/async_await.rs:+1:5: +1:14
+        }
+    }
+    scope 4 {
+        debug __awaitee => (((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2])) as variant#4).0: impl std::future::Future<Output = ()>); // in scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        let _33: ();                     // in scope 4 at $DIR/async_await.rs:+2:5: +2:14
+        scope 5 {
+        }
+        scope 6 {
+            debug result => _33;         // in scope 6 at $DIR/async_await.rs:+2:5: +2:14
+        }
+    }
+
+    bb0: {
+        _39 = discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2]))); // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        switchInt(move _39) -> [0: bb1, 1: bb29, 3: bb27, 4: bb28, otherwise: bb30]; // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    }
+
+    bb1: {
+        _38 = move _2;                   // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        StorageLive(_3);                 // scope 0 at $DIR/async_await.rs:+1:5: +1:14
+        StorageLive(_4);                 // scope 0 at $DIR/async_await.rs:+1:8: +1:14
+        StorageLive(_5);                 // scope 0 at $DIR/async_await.rs:+1:5: +1:8
+        _5 = a() -> bb2;                 // scope 0 at $DIR/async_await.rs:+1:5: +1:8
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:15:5: 15:6
+                                         // + literal: Const { ty: fn() -> impl Future<Output = ()> {a}, val: Value(<ZST>) }
+    }
+
+    bb2: {
+        _4 = <impl Future<Output = ()> as IntoFuture>::into_future(move _5) -> bb3; // scope 0 at $DIR/async_await.rs:+1:8: +1:14
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:15:8: 15:14
+                                         // + literal: Const { ty: fn(impl Future<Output = ()>) -> <impl Future<Output = ()> as IntoFuture>::IntoFuture {<impl Future<Output = ()> as IntoFuture>::into_future}, val: Value(<ZST>) }
+    }
+
+    bb3: {
+        StorageDead(_5);                 // scope 0 at $DIR/async_await.rs:+1:13: +1:14
+        nop;                             // scope 0 at $DIR/async_await.rs:+1:8: +1:14
+        (((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2])) as variant#3).0: impl std::future::Future<Output = ()>) = move _4; // scope 0 at $DIR/async_await.rs:+1:8: +1:14
+        goto -> bb4;                     // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+    }
+
+    bb4: {
+        StorageLive(_8);                 // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        StorageLive(_9);                 // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        StorageLive(_10);                // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+        StorageLive(_11);                // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+        StorageLive(_12);                // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+        _12 = &mut (((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2])) as variant#3).0: impl std::future::Future<Output = ()>); // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+        _11 = &mut (*_12);               // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+        _10 = Pin::<&mut impl Future<Output = ()>>::new_unchecked(move _11) -> bb5; // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:15:8: 15:14
+                                         // + literal: Const { ty: unsafe fn(&mut impl Future<Output = ()>) -> Pin<&mut impl Future<Output = ()>> {Pin::<&mut impl Future<Output = ()>>::new_unchecked}, val: Value(<ZST>) }
+    }
+
+    bb5: {
+        StorageDead(_11);                // scope 2 at $DIR/async_await.rs:+1:13: +1:14
+        StorageLive(_13);                // scope 2 at $DIR/async_await.rs:+1:5: +1:14
+        StorageLive(_14);                // scope 2 at $DIR/async_await.rs:+1:5: +1:14
+        StorageLive(_15);                // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+        _15 = _38;                       // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+        _14 = move _15;                  // scope 2 at $DIR/async_await.rs:+1:5: +1:14
+        goto -> bb6;                     // scope 2 at $DIR/async_await.rs:+1:5: +1:14
+    }
+
+    bb6: {
+        _13 = &mut (*_14);               // scope 2 at $DIR/async_await.rs:+1:5: +1:14
+        StorageDead(_15);                // scope 2 at $DIR/async_await.rs:+1:13: +1:14
+        _9 = <impl Future<Output = ()> as Future>::poll(move _10, move _13) -> bb7; // scope 2 at $DIR/async_await.rs:+1:8: +1:14
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:15:8: 15:14
+                                         // + literal: Const { ty: for<'a, 'b, 'c> fn(Pin<&'a mut impl Future<Output = ()>>, &'b mut Context<'c>) -> Poll<<impl Future<Output = ()> as Future>::Output> {<impl Future<Output = ()> as Future>::poll}, val: Value(<ZST>) }
+    }
+
+    bb7: {
+        StorageDead(_13);                // scope 2 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_10);                // scope 2 at $DIR/async_await.rs:+1:13: +1:14
+        _16 = discriminant(_9);          // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        switchInt(move _16) -> [0: bb10, 1: bb8, otherwise: bb9]; // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+    }
+
+    bb8: {
+        _8 = const ();                   // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        StorageDead(_14);                // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_12);                // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_9);                 // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_8);                 // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageLive(_19);                // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        StorageLive(_20);                // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        _20 = ();                        // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        Deinit(_0);                      // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        discriminant(_0) = 1;            // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2]))) = 3; // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        return;                          // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+    }
+
+    bb9: {
+        unreachable;                     // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+    }
+
+    bb10: {
+        StorageLive(_17);                // scope 1 at $DIR/async_await.rs:+1:5: +1:14
+        _17 = ((_9 as Ready).0: ());     // scope 1 at $DIR/async_await.rs:+1:5: +1:14
+        _3 = _17;                        // scope 3 at $DIR/async_await.rs:+1:5: +1:14
+        StorageDead(_17);                // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_14);                // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_12);                // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_9);                 // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        StorageDead(_8);                 // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        goto -> bb12;                    // scope 0 at $DIR/async_await.rs:+1:13: +1:14
+    }
+
+    bb11: {
+        StorageDead(_20);                // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        _38 = move _19;                  // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        StorageDead(_19);                // scope 1 at $DIR/async_await.rs:+1:13: +1:14
+        _7 = const ();                   // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+        goto -> bb4;                     // scope 1 at $DIR/async_await.rs:+1:8: +1:14
+    }
+
+    bb12: {
+        nop;                             // scope 0 at $DIR/async_await.rs:+1:13: +1:14
+        goto -> bb13;                    // scope 0 at $DIR/async_await.rs:+1:14: +1:15
+    }
+
+    bb13: {
+        StorageDead(_4);                 // scope 0 at $DIR/async_await.rs:+1:14: +1:15
+        StorageDead(_3);                 // scope 0 at $DIR/async_await.rs:+1:14: +1:15
+        StorageLive(_21);                // scope 0 at $DIR/async_await.rs:+2:8: +2:14
+        StorageLive(_22);                // scope 0 at $DIR/async_await.rs:+2:5: +2:8
+        _22 = a() -> bb14;               // scope 0 at $DIR/async_await.rs:+2:5: +2:8
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:16:5: 16:6
+                                         // + literal: Const { ty: fn() -> impl Future<Output = ()> {a}, val: Value(<ZST>) }
+    }
+
+    bb14: {
+        _21 = <impl Future<Output = ()> as IntoFuture>::into_future(move _22) -> bb15; // scope 0 at $DIR/async_await.rs:+2:8: +2:14
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:16:8: 16:14
+                                         // + literal: Const { ty: fn(impl Future<Output = ()>) -> <impl Future<Output = ()> as IntoFuture>::IntoFuture {<impl Future<Output = ()> as IntoFuture>::into_future}, val: Value(<ZST>) }
+    }
+
+    bb15: {
+        StorageDead(_22);                // scope 0 at $DIR/async_await.rs:+2:13: +2:14
+        nop;                             // scope 0 at $DIR/async_await.rs:+2:8: +2:14
+        (((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2])) as variant#4).0: impl std::future::Future<Output = ()>) = move _21; // scope 0 at $DIR/async_await.rs:+2:8: +2:14
+        goto -> bb16;                    // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+    }
+
+    bb16: {
+        StorageLive(_24);                // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        StorageLive(_25);                // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        StorageLive(_26);                // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+        StorageLive(_27);                // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+        StorageLive(_28);                // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+        _28 = &mut (((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2])) as variant#4).0: impl std::future::Future<Output = ()>); // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+        _27 = &mut (*_28);               // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+        _26 = Pin::<&mut impl Future<Output = ()>>::new_unchecked(move _27) -> bb17; // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:16:8: 16:14
+                                         // + literal: Const { ty: unsafe fn(&mut impl Future<Output = ()>) -> Pin<&mut impl Future<Output = ()>> {Pin::<&mut impl Future<Output = ()>>::new_unchecked}, val: Value(<ZST>) }
+    }
+
+    bb17: {
+        StorageDead(_27);                // scope 5 at $DIR/async_await.rs:+2:13: +2:14
+        StorageLive(_29);                // scope 5 at $DIR/async_await.rs:+2:5: +2:14
+        StorageLive(_30);                // scope 5 at $DIR/async_await.rs:+2:5: +2:14
+        StorageLive(_31);                // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+        _31 = _38;                       // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+        _30 = move _31;                  // scope 5 at $DIR/async_await.rs:+2:5: +2:14
+        goto -> bb18;                    // scope 5 at $DIR/async_await.rs:+2:5: +2:14
+    }
+
+    bb18: {
+        _29 = &mut (*_30);               // scope 5 at $DIR/async_await.rs:+2:5: +2:14
+        StorageDead(_31);                // scope 5 at $DIR/async_await.rs:+2:13: +2:14
+        _25 = <impl Future<Output = ()> as Future>::poll(move _26, move _29) -> bb19; // scope 5 at $DIR/async_await.rs:+2:8: +2:14
+                                         // mir::Constant
+                                         // + span: $DIR/async_await.rs:16:8: 16:14
+                                         // + literal: Const { ty: for<'a, 'b, 'c> fn(Pin<&'a mut impl Future<Output = ()>>, &'b mut Context<'c>) -> Poll<<impl Future<Output = ()> as Future>::Output> {<impl Future<Output = ()> as Future>::poll}, val: Value(<ZST>) }
+    }
+
+    bb19: {
+        StorageDead(_29);                // scope 5 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_26);                // scope 5 at $DIR/async_await.rs:+2:13: +2:14
+        _32 = discriminant(_25);         // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        switchInt(move _32) -> [0: bb22, 1: bb20, otherwise: bb21]; // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+    }
+
+    bb20: {
+        _24 = const ();                  // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        StorageDead(_30);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_28);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_25);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_24);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageLive(_35);                // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        StorageLive(_36);                // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        _36 = ();                        // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        Deinit(_0);                      // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        discriminant(_0) = 1;            // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2]))) = 4; // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        return;                          // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+    }
+
+    bb21: {
+        unreachable;                     // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+    }
+
+    bb22: {
+        StorageLive(_33);                // scope 4 at $DIR/async_await.rs:+2:5: +2:14
+        _33 = ((_25 as Ready).0: ());    // scope 4 at $DIR/async_await.rs:+2:5: +2:14
+        _37 = _33;                       // scope 6 at $DIR/async_await.rs:+2:5: +2:14
+        StorageDead(_33);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_30);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_28);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_25);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        StorageDead(_24);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        goto -> bb24;                    // scope 0 at $DIR/async_await.rs:+2:13: +2:14
+    }
+
+    bb23: {
+        StorageDead(_36);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        _38 = move _35;                  // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        StorageDead(_35);                // scope 4 at $DIR/async_await.rs:+2:13: +2:14
+        _7 = const ();                   // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+        goto -> bb16;                    // scope 4 at $DIR/async_await.rs:+2:8: +2:14
+    }
+
+    bb24: {
+        nop;                             // scope 0 at $DIR/async_await.rs:+2:13: +2:14
+        goto -> bb25;                    // scope 0 at $DIR/async_await.rs:+3:1: +3:2
+    }
+
+    bb25: {
+        StorageDead(_21);                // scope 0 at $DIR/async_await.rs:+3:1: +3:2
+        goto -> bb26;                    // scope 0 at $DIR/async_await.rs:+3:1: +3:2
+    }
+
+    bb26: {
+        Deinit(_0);                      // scope 0 at $DIR/async_await.rs:+3:2: +3:2
+        ((_0 as Ready).0: ()) = move _37; // scope 0 at $DIR/async_await.rs:+3:2: +3:2
+        discriminant(_0) = 0;            // scope 0 at $DIR/async_await.rs:+3:2: +3:2
+        discriminant((*(_1.0: &mut [async fn body@$DIR/async_await.rs:14:18: 17:2]))) = 1; // scope 0 at $DIR/async_await.rs:+3:2: +3:2
+        return;                          // scope 0 at $DIR/async_await.rs:+3:2: +3:2
+    }
+
+    bb27: {
+        StorageLive(_3);                 // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        StorageLive(_4);                 // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        StorageLive(_19);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        StorageLive(_20);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        _19 = move _2;                   // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        goto -> bb11;                    // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    }
+
+    bb28: {
+        StorageLive(_21);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        StorageLive(_35);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        StorageLive(_36);                // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        _35 = move _2;                   // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+        goto -> bb23;                    // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    }
+
+    bb29: {
+        assert(const false, "`async fn` resumed after completion") -> bb29; // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    }
+
+    bb30: {
+        unreachable;                     // scope 0 at $DIR/async_await.rs:+0:18: +3:2
+    }
+}

tests/mir-opt/building/async_await.rs

@@ -0,0 +1,17 @@
+// This test makes sure that the generator MIR pass eliminates all calls to
+// `get_context`, and that the MIR argument type for an async fn and all locals
+// related to `yield` are `&mut Context`, and its return type is `Poll`.
+
+// edition:2018
+// compile-flags: -C panic=abort
+
+#![crate_type = "lib"]
+
+// EMIT_MIR async_await.a-{closure#0}.generator_resume.0.mir
+async fn a() {}
+
+// EMIT_MIR async_await.b-{closure#0}.generator_resume.0.mir
+pub async fn b() {
+    a().await;
+    a().await
+}