Fix drop shim for AsyncFnOnce closure, AsyncFnMut shim for AsyncFn closure

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -546,7 +546,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             | ty::InstanceDef::ReifyShim(..)
             | ty::InstanceDef::ClosureOnceShim { .. }
             | ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
-            | ty::InstanceDef::CoroutineByMoveShim { .. }
+            | ty::InstanceDef::CoroutineKindShim { .. }
             | ty::InstanceDef::FnPtrShim(..)
             | ty::InstanceDef::DropGlue(..)
             | ty::InstanceDef::CloneShim(..)

compiler/rustc_middle/src/mir/mod.rs

@@ -265,7 +265,7 @@ pub struct CoroutineInfo<'tcx> {
     /// The body of the coroutine, modified to take its upvars by move rather than by ref.
     ///
     /// This is used by coroutine-closures, which must return a different flavor of coroutine
-    /// when called using `AsyncFnOnce::call_once`. It is produced by the `ByMoveBody` which
+    /// when called using `AsyncFnOnce::call_once`. It is produced by the `ByMoveBody` pass which
     /// is run right after building the initial MIR, and will only be populated for coroutines
     /// which come out of the async closure desugaring.
     ///
@@ -274,6 +274,13 @@ pub struct CoroutineInfo<'tcx> {
     /// using `run_passes`.
     pub by_move_body: Option<Body<'tcx>>,
 
+    /// The body of the coroutine, modified to take its upvars by mutable ref rather than by
+    /// immutable ref.
+    ///
+    /// FIXME(async_closures): This is literally the same body as the parent body. Find a better
+    /// way to represent the by-mut signature (or cap the closure-kind of the coroutine).
+    pub by_mut_body: Option<Body<'tcx>>,
+
     /// The layout of a coroutine. This field is populated after the state transform pass.
     pub coroutine_layout: Option<CoroutineLayout<'tcx>>,
 
@@ -294,6 +301,7 @@ impl<'tcx> CoroutineInfo<'tcx> {
             yield_ty: Some(yield_ty),
             resume_ty: Some(resume_ty),
             by_move_body: None,
+            by_mut_body: None,
             coroutine_drop: None,
             coroutine_layout: None,
         }
@@ -604,6 +612,14 @@ impl<'tcx> Body<'tcx> {
         self.coroutine.as_ref().and_then(|coroutine| coroutine.coroutine_drop.as_ref())
     }
 
+    pub fn coroutine_by_move_body(&self) -> Option<&Body<'tcx>> {
+        self.coroutine.as_ref()?.by_move_body.as_ref()
+    }
+
+    pub fn coroutine_by_mut_body(&self) -> Option<&Body<'tcx>> {
+        self.coroutine.as_ref()?.by_mut_body.as_ref()
+    }
+
     #[inline]
     pub fn coroutine_kind(&self) -> Option<CoroutineKind> {
         self.coroutine.as_ref().map(|coroutine| coroutine.coroutine_kind)

compiler/rustc_middle/src/mir/mono.rs

@@ -403,7 +403,7 @@ impl<'tcx> CodegenUnit<'tcx> {
                             | InstanceDef::Virtual(..)
                             | InstanceDef::ClosureOnceShim { .. }
                             | InstanceDef::ConstructCoroutineInClosureShim { .. }
-                            | InstanceDef::CoroutineByMoveShim { .. }
+                            | InstanceDef::CoroutineKindShim { .. }
                             | InstanceDef::DropGlue(..)
                             | InstanceDef::CloneShim(..)
                             | InstanceDef::ThreadLocalShim(..)

compiler/rustc_middle/src/mir/visit.rs

@@ -346,7 +346,7 @@ macro_rules! make_mir_visitor {
                         ty::InstanceDef::ThreadLocalShim(_def_id) |
                         ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |
                         ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id: _def_id, target_kind: _ } |
-                        ty::InstanceDef::CoroutineByMoveShim { coroutine_def_id: _def_id } |
+                        ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id, target_kind: _ } |
                         ty::InstanceDef::DropGlue(_def_id, None) => {}
 
                         ty::InstanceDef::FnPtrShim(_def_id, ty) |

compiler/rustc_middle/src/ty/instance.rs

@@ -102,10 +102,12 @@ pub enum InstanceDef<'tcx> {
     },
 
     /// `<[coroutine] as Future>::poll`, but for coroutines produced when `AsyncFnOnce`
-    /// is called on a coroutine-closure whose closure kind is not `FnOnce`. This
+    /// is called on a coroutine-closure whose closure kind greater than `FnOnce`, or
-    /// will select the body that is produced by the `ByMoveBody` transform, and thus
+    /// similarly for `AsyncFnMut`.
+    ///
+    /// This will select the body that is produced by the `ByMoveBody` transform, and thus
     /// take and use all of its upvars by-move rather than by-ref.
-    CoroutineByMoveShim { coroutine_def_id: DefId },
+    CoroutineKindShim { coroutine_def_id: DefId, target_kind: ty::ClosureKind },
 
     /// Compiler-generated accessor for thread locals which returns a reference to the thread local
     /// the `DefId` defines. This is used to export thread locals from dylibs on platforms lacking
@@ -192,7 +194,7 @@ impl<'tcx> InstanceDef<'tcx> {
                 coroutine_closure_def_id: def_id,
                 target_kind: _,
             }
-            | ty::InstanceDef::CoroutineByMoveShim { coroutine_def_id: def_id }
+            | ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id, target_kind: _ }
             | InstanceDef::DropGlue(def_id, _)
             | InstanceDef::CloneShim(def_id, _)
             | InstanceDef::FnPtrAddrShim(def_id, _) => def_id,
@@ -213,7 +215,7 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::Intrinsic(..)
             | InstanceDef::ClosureOnceShim { .. }
             | ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
-            | ty::InstanceDef::CoroutineByMoveShim { .. }
+            | ty::InstanceDef::CoroutineKindShim { .. }
             | InstanceDef::DropGlue(..)
             | InstanceDef::CloneShim(..)
             | InstanceDef::FnPtrAddrShim(..) => None,
@@ -310,7 +312,7 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::DropGlue(_, Some(_)) => false,
             InstanceDef::ClosureOnceShim { .. }
             | InstanceDef::ConstructCoroutineInClosureShim { .. }
-            | InstanceDef::CoroutineByMoveShim { .. }
+            | InstanceDef::CoroutineKindShim { .. }
             | InstanceDef::DropGlue(..)
             | InstanceDef::Item(_)
             | InstanceDef::Intrinsic(..)
@@ -349,7 +351,7 @@ fn fmt_instance(
         InstanceDef::FnPtrShim(_, ty) => write!(f, " - shim({ty})"),
         InstanceDef::ClosureOnceShim { .. } => write!(f, " - shim"),
         InstanceDef::ConstructCoroutineInClosureShim { .. } => write!(f, " - shim"),
-        InstanceDef::CoroutineByMoveShim { .. } => write!(f, " - shim"),
+        InstanceDef::CoroutineKindShim { .. } => write!(f, " - shim"),
         InstanceDef::DropGlue(_, None) => write!(f, " - shim(None)"),
         InstanceDef::DropGlue(_, Some(ty)) => write!(f, " - shim(Some({ty}))"),
         InstanceDef::CloneShim(_, ty) => write!(f, " - shim({ty})"),
@@ -651,13 +653,11 @@ impl<'tcx> Instance<'tcx> {
             if args.as_coroutine().kind_ty() == id_args.as_coroutine().kind_ty() {
                 Some(Instance { def: ty::InstanceDef::Item(coroutine_def_id), args })
             } else {
-                assert_eq!(
-                    args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
-                    ty::ClosureKind::FnOnce,
-                    "FIXME(async_closures): Generate a by-mut body here."
-                );
                 Some(Instance {
-                    def: ty::InstanceDef::CoroutineByMoveShim { coroutine_def_id },
+                    def: ty::InstanceDef::CoroutineKindShim {
+                        coroutine_def_id,
+                        target_kind: args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
+                    },
                     args,
                 })
             }

compiler/rustc_middle/src/ty/mod.rs

@@ -1681,7 +1681,7 @@ impl<'tcx> TyCtxt<'tcx> {
             | ty::InstanceDef::Virtual(..)
             | ty::InstanceDef::ClosureOnceShim { .. }
             | ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
-            | ty::InstanceDef::CoroutineByMoveShim { .. }
+            | ty::InstanceDef::CoroutineKindShim { .. }
             | ty::InstanceDef::DropGlue(..)
             | ty::InstanceDef::CloneShim(..)
             | ty::InstanceDef::ThreadLocalShim(..)

compiler/rustc_middle/src/ty/print/mod.rs

@@ -3,6 +3,7 @@ use crate::ty::{self, Ty, TyCtxt};
 
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::sso::SsoHashSet;
+use rustc_hir as hir;
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId};
 use rustc_hir::definitions::{DefPathData, DisambiguatedDefPathData};
 
@@ -130,8 +131,24 @@ pub trait Printer<'tcx>: Sized {
                     parent_args = &args[..generics.parent_count.min(args.len())];
 
                     match key.disambiguated_data.data {
-                        // Closures' own generics are only captures, don't print them.
+                        DefPathData::Closure => {
-                        DefPathData::Closure => {}
+                            // FIXME(async_closures): This is somewhat ugly.
+                            // We need to additionally print the `kind` field of a closure if
+                            // it is desugared from a coroutine-closure.
+                            if let Some(hir::CoroutineKind::Desugared(
+                                _,
+                                hir::CoroutineSource::Closure,
+                            )) = self.tcx().coroutine_kind(def_id)
+                                && args.len() >= parent_args.len() + 1
+                            {
+                                return self.path_generic_args(
+                                    |cx| cx.print_def_path(def_id, parent_args),
+                                    &args[..parent_args.len() + 1][..1],
+                                );
+                            } else {
+                                // Closures' own generics are only captures, don't print them.
+                            }
+                        }
                         // This covers both `DefKind::AnonConst` and `DefKind::InlineConst`.
                         // Anon consts doesn't have their own generics, and inline consts' own
                         // generics are their inferred types, so don't print them.

compiler/rustc_mir_transform/src/coroutine/by_move_body.rs

@@ -6,7 +6,7 @@
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_hir as hir;
 use rustc_middle::mir::visit::MutVisitor;
-use rustc_middle::mir::{self, MirPass};
+use rustc_middle::mir::{self, dump_mir, MirPass};
 use rustc_middle::ty::{self, InstanceDef, Ty, TyCtxt};
 use rustc_target::abi::FieldIdx;
 
@@ -24,7 +24,9 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
         };
         let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
         let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!() };
-        if args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap() == ty::ClosureKind::FnOnce {
+
+        let coroutine_kind = args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap();
+        if coroutine_kind == ty::ClosureKind::FnOnce {
             return;
         }
 
@@ -58,14 +60,49 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
 
         let mut by_move_body = body.clone();
         MakeByMoveBody { tcx, by_ref_fields, by_move_coroutine_ty }.visit_body(&mut by_move_body);
+        dump_mir(tcx, false, "coroutine_by_move", &0, &by_move_body, |_, _| Ok(()));
         by_move_body.source = mir::MirSource {
-            instance: InstanceDef::CoroutineByMoveShim {
+            instance: InstanceDef::CoroutineKindShim {
                 coroutine_def_id: coroutine_def_id.to_def_id(),
+                target_kind: ty::ClosureKind::FnOnce,
             },
             promoted: None,
         };
-
         body.coroutine.as_mut().unwrap().by_move_body = Some(by_move_body);
+
+        // If this is coming from an `AsyncFn` coroutine-closure, we must also create a by-mut body.
+        // This is actually just a copy of the by-ref body, but with a different self type.
+        // FIXME(async_closures): We could probably unify this with the by-ref body somehow.
+        if coroutine_kind == ty::ClosureKind::Fn {
+            let by_mut_coroutine_ty = Ty::new_coroutine(
+                tcx,
+                coroutine_def_id.to_def_id(),
+                ty::CoroutineArgs::new(
+                    tcx,
+                    ty::CoroutineArgsParts {
+                        parent_args: args.as_coroutine().parent_args(),
+                        kind_ty: Ty::from_closure_kind(tcx, ty::ClosureKind::FnMut),
+                        resume_ty: args.as_coroutine().resume_ty(),
+                        yield_ty: args.as_coroutine().yield_ty(),
+                        return_ty: args.as_coroutine().return_ty(),
+                        witness: args.as_coroutine().witness(),
+                        tupled_upvars_ty: args.as_coroutine().tupled_upvars_ty(),
+                    },
+                )
+                .args,
+            );
+            let mut by_mut_body = body.clone();
+            by_mut_body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty = by_mut_coroutine_ty;
+            dump_mir(tcx, false, "coroutine_by_mut", &0, &by_mut_body, |_, _| Ok(()));
+            by_mut_body.source = mir::MirSource {
+                instance: InstanceDef::CoroutineKindShim {
+                    coroutine_def_id: coroutine_def_id.to_def_id(),
+                    target_kind: ty::ClosureKind::FnMut,
+                },
+                promoted: None,
+            };
+            body.coroutine.as_mut().unwrap().by_mut_body = Some(by_mut_body);
+        }
     }
 }
 

compiler/rustc_mir_transform/src/inline.rs

@@ -318,7 +318,7 @@ impl<'tcx> Inliner<'tcx> {
             | InstanceDef::FnPtrShim(..)
             | InstanceDef::ClosureOnceShim { .. }
             | InstanceDef::ConstructCoroutineInClosureShim { .. }
-            | InstanceDef::CoroutineByMoveShim { .. }
+            | InstanceDef::CoroutineKindShim { .. }
             | InstanceDef::DropGlue(..)
             | InstanceDef::CloneShim(..)
             | InstanceDef::ThreadLocalShim(..)

compiler/rustc_mir_transform/src/inline/cycle.rs

@@ -88,7 +88,7 @@ pub(crate) fn mir_callgraph_reachable<'tcx>(
                 | InstanceDef::FnPtrShim(..)
                 | InstanceDef::ClosureOnceShim { .. }
                 | InstanceDef::ConstructCoroutineInClosureShim { .. }
-                | InstanceDef::CoroutineByMoveShim { .. }
+                | InstanceDef::CoroutineKindShim { .. }
                 | InstanceDef::ThreadLocalShim { .. }
                 | InstanceDef::CloneShim(..) => {}
 

compiler/rustc_mir_transform/src/pass_manager.rs

@@ -190,10 +190,13 @@ fn run_passes_inner<'tcx>(
         body.pass_count = 1;
     }
 
-    if let Some(coroutine) = body.coroutine.as_mut()
+    if let Some(coroutine) = body.coroutine.as_mut() {
-        && let Some(by_move_body) = coroutine.by_move_body.as_mut()
+        if let Some(by_move_body) = coroutine.by_move_body.as_mut() {
-    {
+            run_passes_inner(tcx, by_move_body, passes, phase_change, validate_each);
-        run_passes_inner(tcx, by_move_body, passes, phase_change, validate_each);
+        }
+        if let Some(by_mut_body) = coroutine.by_mut_body.as_mut() {
+            run_passes_inner(tcx, by_mut_body, passes, phase_change, validate_each);
+        }
     }
 }
 

compiler/rustc_mir_transform/src/shim.rs

@@ -72,32 +72,70 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
         } => match target_kind {
             ty::ClosureKind::Fn => unreachable!("shouldn't be building shim for Fn"),
             ty::ClosureKind::FnMut => {
-                let body = build_construct_coroutine_by_mut_shim(tcx, coroutine_closure_def_id);
+                // No need to optimize the body, it has already been optimized
-                // No need to optimize the body, it has already been optimized.
+                // since we steal it from the `AsyncFn::call` body and just fix
-                return body;
+                // the return type.
+                return build_construct_coroutine_by_mut_shim(tcx, coroutine_closure_def_id);
             }
             ty::ClosureKind::FnOnce => {
                 build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id)
             }
         },
 
-        ty::InstanceDef::CoroutineByMoveShim { coroutine_def_id } => {
+        ty::InstanceDef::CoroutineKindShim { coroutine_def_id, target_kind } => match target_kind {
-            return tcx
+            ty::ClosureKind::Fn => unreachable!(),
-                .optimized_mir(coroutine_def_id)
+            ty::ClosureKind::FnMut => {
-                .coroutine
+                return tcx
-                .as_ref()
+                    .optimized_mir(coroutine_def_id)
-                .unwrap()
+                    .coroutine_by_mut_body()
-                .by_move_body
+                    .unwrap()
-                .as_ref()
+                    .clone();
-                .unwrap()
+            }
-                .clone();
+            ty::ClosureKind::FnOnce => {
-        }
+                return tcx
+                    .optimized_mir(coroutine_def_id)
+                    .coroutine_by_move_body()
+                    .unwrap()
+                    .clone();
+            }
+        },
 
         ty::InstanceDef::DropGlue(def_id, ty) => {
             // FIXME(#91576): Drop shims for coroutines aren't subject to the MIR passes at the end
             // of this function. Is this intentional?
             if let Some(ty::Coroutine(coroutine_def_id, args)) = ty.map(Ty::kind) {
-                let body = tcx.optimized_mir(*coroutine_def_id).coroutine_drop().unwrap();
+                let coroutine_body = tcx.optimized_mir(*coroutine_def_id);
+
+                let ty::Coroutine(_, id_args) = *tcx.type_of(coroutine_def_id).skip_binder().kind()
+                else {
+                    bug!()
+                };
+
+                // If this is a regular coroutine, grab its drop shim. If this is a coroutine
+                // that comes from a coroutine-closure, and the kind ty differs from the "maximum"
+                // kind that it supports, then grab the appropriate drop shim. This ensures that
+                // the future returned by `<[coroutine-closure] as AsyncFnOnce>::call_once` will
+                // drop the coroutine-closure's upvars.
+                let body = if id_args.as_coroutine().kind_ty() == args.as_coroutine().kind_ty() {
+                    coroutine_body.coroutine_drop().unwrap()
+                } else {
+                    match args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap() {
+                        ty::ClosureKind::Fn => {
+                            unreachable!()
+                        }
+                        ty::ClosureKind::FnMut => coroutine_body
+                            .coroutine_by_mut_body()
+                            .unwrap()
+                            .coroutine_drop()
+                            .unwrap(),
+                        ty::ClosureKind::FnOnce => coroutine_body
+                            .coroutine_by_move_body()
+                            .unwrap()
+                            .coroutine_drop()
+                            .unwrap(),
+                    }
+                };
+
                 let mut body = EarlyBinder::bind(body.clone()).instantiate(tcx, args);
                 debug!("make_shim({:?}) = {:?}", instance, body);
 
@@ -1076,7 +1114,11 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
         target_kind: ty::ClosureKind::FnOnce,
     });
 
-    new_body(source, IndexVec::from_elem_n(start_block, 1), locals, sig.inputs().len(), span)
+    let body =
+        new_body(source, IndexVec::from_elem_n(start_block, 1), locals, sig.inputs().len(), span);
+    dump_mir(tcx, false, "coroutine_closure_by_move", &0, &body, |_, _| Ok(()));
+
+    body
 }
 
 fn build_construct_coroutine_by_mut_shim<'tcx>(
@@ -1110,5 +1152,8 @@ fn build_construct_coroutine_by_mut_shim<'tcx>(
         target_kind: ty::ClosureKind::FnMut,
     });
 
+    body.pass_count = 0;
+    dump_mir(tcx, false, "coroutine_closure_by_mut", &0, &body, |_, _| Ok(()));
+
     body
 }

compiler/rustc_monomorphize/src/collector.rs

@@ -984,7 +984,7 @@ fn visit_instance_use<'tcx>(
         | ty::InstanceDef::ReifyShim(..)
         | ty::InstanceDef::ClosureOnceShim { .. }
         | ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
-        | ty::InstanceDef::CoroutineByMoveShim { .. }
+        | ty::InstanceDef::CoroutineKindShim { .. }
         | ty::InstanceDef::Item(..)
         | ty::InstanceDef::FnPtrShim(..)
         | ty::InstanceDef::CloneShim(..)

compiler/rustc_monomorphize/src/partitioning.rs

@@ -621,7 +621,7 @@ fn characteristic_def_id_of_mono_item<'tcx>(
                 | ty::InstanceDef::FnPtrShim(..)
                 | ty::InstanceDef::ClosureOnceShim { .. }
                 | ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
-                | ty::InstanceDef::CoroutineByMoveShim { .. }
+                | ty::InstanceDef::CoroutineKindShim { .. }
                 | ty::InstanceDef::Intrinsic(..)
                 | ty::InstanceDef::DropGlue(..)
                 | ty::InstanceDef::Virtual(..)
@@ -786,7 +786,7 @@ fn mono_item_visibility<'tcx>(
         | InstanceDef::Intrinsic(..)
         | InstanceDef::ClosureOnceShim { .. }
         | InstanceDef::ConstructCoroutineInClosureShim { .. }
-        | InstanceDef::CoroutineByMoveShim { .. }
+        | InstanceDef::CoroutineKindShim { .. }
         | InstanceDef::DropGlue(..)
         | InstanceDef::CloneShim(..)
         | InstanceDef::FnPtrAddrShim(..) => return Visibility::Hidden,

compiler/rustc_smir/src/rustc_smir/convert/ty.rs

@@ -800,7 +800,7 @@ impl<'tcx> Stable<'tcx> for ty::Instance<'tcx> {
             | ty::InstanceDef::FnPtrAddrShim(..)
             | ty::InstanceDef::ClosureOnceShim { .. }
             | ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
-            | ty::InstanceDef::CoroutineByMoveShim { .. }
+            | ty::InstanceDef::CoroutineKindShim { .. }
             | ty::InstanceDef::ThreadLocalShim(..)
             | ty::InstanceDef::DropGlue(..)
             | ty::InstanceDef::CloneShim(..)

compiler/rustc_symbol_mangling/src/legacy.rs

@@ -64,16 +64,29 @@ pub(super) fn mangle<'tcx>(
         )
         .unwrap();
 
-    if let ty::InstanceDef::ThreadLocalShim(..) = instance.def {
+    match instance.def {
-        let _ = printer.write_str("{{tls-shim}}");
+        ty::InstanceDef::ThreadLocalShim(..) => {
-    }
+            printer.write_str("{{tls-shim}}").unwrap();
-
+        }
-    if let ty::InstanceDef::VTableShim(..) = instance.def {
+        ty::InstanceDef::VTableShim(..) => {
-        let _ = printer.write_str("{{vtable-shim}}");
+            printer.write_str("{{vtable-shim}}").unwrap();
-    }
+        }
-
+        ty::InstanceDef::ReifyShim(..) => {
-    if let ty::InstanceDef::ReifyShim(..) = instance.def {
+            printer.write_str("{{reify-shim}}").unwrap();
-        let _ = printer.write_str("{{reify-shim}}");
+        }
+        // FIXME(async_closures): This shouldn't be needed when we fix
+        // `Instance::ty`/`Instance::def_id`.
+        ty::InstanceDef::ConstructCoroutineInClosureShim { target_kind, .. }
+        | ty::InstanceDef::CoroutineKindShim { target_kind, .. } => match target_kind {
+            ty::ClosureKind::Fn => unreachable!(),
+            ty::ClosureKind::FnMut => {
+                printer.write_str("{{fn-mut-shim}}").unwrap();
+            }
+            ty::ClosureKind::FnOnce => {
+                printer.write_str("{{fn-once-shim}}").unwrap();
+            }
+        },
+        _ => {}
     }
 
     printer.path.finish(hash)

compiler/rustc_symbol_mangling/src/v0.rs

@@ -46,6 +46,13 @@ pub(super) fn mangle<'tcx>(
         ty::InstanceDef::VTableShim(_) => Some("vtable"),
         ty::InstanceDef::ReifyShim(_) => Some("reify"),
 
+        ty::InstanceDef::ConstructCoroutineInClosureShim { target_kind, .. }
+        | ty::InstanceDef::CoroutineKindShim { target_kind, .. } => match target_kind {
+            ty::ClosureKind::Fn => unreachable!(),
+            ty::ClosureKind::FnMut => Some("fn_mut"),
+            ty::ClosureKind::FnOnce => Some("fn_once"),
+        },
+
         _ => None,
     };
 

compiler/rustc_ty_utils/src/abi.rs

@@ -145,7 +145,7 @@ fn fn_sig_for_fn_abi<'tcx>(
             )
         }
         ty::Coroutine(did, args) => {
-            // FIXME(async_closures): This isn't right for `CoroutineByMoveShim`.
+            // FIXME(async_closures): This isn't right for `CoroutineKindShim`.
 
             let coroutine_kind = tcx.coroutine_kind(did).unwrap();
             let sig = args.as_coroutine().sig();

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_move.0.panic-abort.mir

@@ -0,0 +1,47 @@
+// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_move
+
+fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10}, _2: ResumeTy) -> ()
+yields ()
+ {
+    debug _task_context => _2;
+    debug a => (_1.0: i32);
+    debug b => (_1.1: i32);
+    let mut _0: ();
+    let _3: i32;
+    scope 1 {
+        debug a => _3;
+        let _4: &i32;
+        scope 2 {
+            debug a => _4;
+            let _5: &i32;
+            scope 3 {
+                debug b => _5;
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_3);
+        _3 = (_1.0: i32);
+        FakeRead(ForLet(None), _3);
+        StorageLive(_4);
+        _4 = &_3;
+        FakeRead(ForLet(None), _4);
+        StorageLive(_5);
+        _5 = &(_1.1: i32);
+        FakeRead(ForLet(None), _5);
+        _0 = const ();
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        drop(_1) -> [return: bb1, unwind: bb2];
+    }
+
+    bb1: {
+        return;
+    }
+
+    bb2 (cleanup): {
+        resume;
+    }
+}

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_move.0.panic-unwind.mir

@@ -0,0 +1,47 @@
+// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_move
+
+fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10}, _2: ResumeTy) -> ()
+yields ()
+ {
+    debug _task_context => _2;
+    debug a => (_1.0: i32);
+    debug b => (_1.1: i32);
+    let mut _0: ();
+    let _3: i32;
+    scope 1 {
+        debug a => _3;
+        let _4: &i32;
+        scope 2 {
+            debug a => _4;
+            let _5: &i32;
+            scope 3 {
+                debug b => _5;
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_3);
+        _3 = (_1.0: i32);
+        FakeRead(ForLet(None), _3);
+        StorageLive(_4);
+        _4 = &_3;
+        FakeRead(ForLet(None), _4);
+        StorageLive(_5);
+        _5 = &(_1.1: i32);
+        FakeRead(ForLet(None), _5);
+        _0 = const ();
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        drop(_1) -> [return: bb1, unwind: bb2];
+    }
+
+    bb1: {
+        return;
+    }
+
+    bb2 (cleanup): {
+        resume;
+    }
+}

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_mut.0.panic-abort.mir

@@ -0,0 +1,47 @@
+// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_mut
+
+fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10}, _2: ResumeTy) -> ()
+yields ()
+ {
+    debug _task_context => _2;
+    debug a => (_1.0: i32);
+    debug b => (*(_1.1: &i32));
+    let mut _0: ();
+    let _3: i32;
+    scope 1 {
+        debug a => _3;
+        let _4: &i32;
+        scope 2 {
+            debug a => _4;
+            let _5: &i32;
+            scope 3 {
+                debug b => _5;
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_3);
+        _3 = (_1.0: i32);
+        FakeRead(ForLet(None), _3);
+        StorageLive(_4);
+        _4 = &_3;
+        FakeRead(ForLet(None), _4);
+        StorageLive(_5);
+        _5 = &(*(_1.1: &i32));
+        FakeRead(ForLet(None), _5);
+        _0 = const ();
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        drop(_1) -> [return: bb1, unwind: bb2];
+    }
+
+    bb1: {
+        return;
+    }
+
+    bb2 (cleanup): {
+        resume;
+    }
+}

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_mut.0.panic-unwind.mir

@@ -0,0 +1,47 @@
+// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_mut
+
+fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10}, _2: ResumeTy) -> ()
+yields ()
+ {
+    debug _task_context => _2;
+    debug a => (_1.0: i32);
+    debug b => (*(_1.1: &i32));
+    let mut _0: ();
+    let _3: i32;
+    scope 1 {
+        debug a => _3;
+        let _4: &i32;
+        scope 2 {
+            debug a => _4;
+            let _5: &i32;
+            scope 3 {
+                debug b => _5;
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_3);
+        _3 = (_1.0: i32);
+        FakeRead(ForLet(None), _3);
+        StorageLive(_4);
+        _4 = &_3;
+        FakeRead(ForLet(None), _4);
+        StorageLive(_5);
+        _5 = &(*(_1.1: &i32));
+        FakeRead(ForLet(None), _5);
+        _0 = const ();
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        drop(_1) -> [return: bb1, unwind: bb2];
+    }
+
+    bb1: {
+        return;
+    }
+
+    bb2 (cleanup): {
+        resume;
+    }
+}

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_move.0.panic-abort.mir

@@ -0,0 +1,10 @@
+// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_move
+
+fn main::{closure#0}::{closure#0}(_1: {coroutine-closure@$DIR/async_closure_shims.rs:39:33: 39:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10} {
+    let mut _0: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10};
+
+    bb0: {
+        _0 = {coroutine@$DIR/async_closure_shims.rs:39:53: 42:10 (#0)} { a: move _2, b: move (_1.0: i32) };
+        return;
+    }
+}

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_move.0.panic-unwind.mir

@@ -0,0 +1,10 @@
+// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_move
+
+fn main::{closure#0}::{closure#0}(_1: {coroutine-closure@$DIR/async_closure_shims.rs:39:33: 39:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10} {
+    let mut _0: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10};
+
+    bb0: {
+        _0 = {coroutine@$DIR/async_closure_shims.rs:39:53: 42:10 (#0)} { a: move _2, b: move (_1.0: i32) };
+        return;
+    }
+}

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_mut.0.panic-abort.mir

@@ -0,0 +1,16 @@
+// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_mut
+
+fn main::{closure#0}::{closure#0}(_1: &mut {coroutine-closure@$DIR/async_closure_shims.rs:39:33: 39:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10} {
+    debug a => _2;
+    debug b => ((*_1).0: i32);
+    let mut _0: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10};
+    let mut _3: &i32;
+
+    bb0: {
+        StorageLive(_3);
+        _3 = &((*_1).0: i32);
+        _0 = {coroutine@$DIR/async_closure_shims.rs:39:53: 42:10 (#0)} { a: _2, b: move _3 };
+        StorageDead(_3);
+        return;
+    }
+}

tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_mut.0.panic-unwind.mir

@@ -0,0 +1,16 @@
+// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_mut
+
+fn main::{closure#0}::{closure#0}(_1: &mut {coroutine-closure@$DIR/async_closure_shims.rs:39:33: 39:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10} {
+    debug a => _2;
+    debug b => ((*_1).0: i32);
+    let mut _0: {async closure body@$DIR/async_closure_shims.rs:39:53: 42:10};
+    let mut _3: &i32;
+
+    bb0: {
+        StorageLive(_3);
+        _3 = &((*_1).0: i32);
+        _0 = {coroutine@$DIR/async_closure_shims.rs:39:53: 42:10 (#0)} { a: _2, b: move _3 };
+        StorageDead(_3);
+        return;
+    }
+}

tests/mir-opt/async_closure_shims.rs

@@ -0,0 +1,46 @@
+// edition:2021
+// skip-filecheck
+// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
+
+#![feature(async_closure, noop_waker, async_fn_traits)]
+
+use std::future::Future;
+use std::ops::{AsyncFnMut, AsyncFnOnce};
+use std::pin::pin;
+use std::task::*;
+
+pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
+    let mut fut = pin!(fut);
+    let ctx = &mut Context::from_waker(Waker::noop());
+
+    loop {
+        match fut.as_mut().poll(ctx) {
+            Poll::Pending => {}
+            Poll::Ready(t) => break t,
+        }
+    }
+}
+
+async fn call_mut(f: &mut impl AsyncFnMut(i32)) {
+    f(0).await;
+}
+
+async fn call_once(f: impl AsyncFnOnce(i32)) {
+    f(1).await;
+}
+
+// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_move.0.mir
+// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_mut.0.mir
+// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_mut.0.mir
+// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_move.0.mir
+fn main() {
+    block_on(async {
+        let b = 2i32;
+        let mut async_closure = async move |a: i32| {
+            let a = &a;
+            let b = &b;
+        };
+        call_mut(&mut async_closure).await;
+        call_once(async_closure).await;
+    });
+}

tests/ui/async-await/async-closures/async-fn-mut-for-async-fn.rs

@@ -0,0 +1,23 @@
+// aux-build:block-on.rs
+// edition:2021
+// run-pass
+
+// FIXME(async_closures): When `fn_sig_for_fn_abi` is fixed, remove this.
+// ignore-pass (test emits codegen-time warnings)
+
+#![feature(async_closure, async_fn_traits)]
+
+extern crate block_on;
+
+use std::ops::AsyncFnMut;
+
+fn main() {
+    block_on::block_on(async {
+        let x = async || {};
+
+        async fn needs_async_fn_mut(mut x: impl AsyncFnMut()) {
+            x().await;
+        }
+        needs_async_fn_mut(x).await;
+    });
+}

tests/ui/async-await/async-closures/async-fn-mut-for-async-fn.stderr

@@ -0,0 +1 @@
+WARN rustc_codegen_ssa::mir::locals Unexpected initial operand type: expected std::pin::Pin<&ReErased mut Coroutine(DefId(0:8 ~ async_fn_mut_for_async_fn[3241]::main::{closure#0}::{closure#0}::{closure#0}), [i16, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_mut_for_async_fn[3241]::main::{closure#0}::{closure#0}::{closure#0}), []), ()])>, found std::pin::Pin<&ReErased mut Coroutine(DefId(0:8 ~ async_fn_mut_for_async_fn[3241]::main::{closure#0}::{closure#0}::{closure#0}), [i8, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_mut_for_async_fn[3241]::main::{closure#0}::{closure#0}::{closure#0}), []), ()])>.See <https://github.com/rust-lang/rust/issues/114858>.

tests/ui/async-await/async-closures/async-fn-once-for-async-fn.stderr

@@ -1,2 +1 @@
-WARN rustc_codegen_ssa::mir::locals Unexpected initial operand type: expected std::pin::Pin<&ReErased mut Coroutine(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), [i32, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), []), ()])>, found std::pin::Pin<&ReErased mut Coroutine(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), [i8, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), []), ()])>.See <https://github.com/rust-lang/rust/issues/114858>.
+WARN rustc_codegen_ssa::mir::locals Unexpected initial operand type: expected std::pin::Pin<&ReErased mut Coroutine(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#0}::{closure#0}), [i32, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#0}::{closure#0}), []), ()])>, found std::pin::Pin<&ReErased mut Coroutine(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#0}::{closure#0}), [i8, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#0}::{closure#0}), []), ()])>.See <https://github.com/rust-lang/rust/issues/114858>.
-WARN rustc_codegen_ssa::mir::locals Unexpected initial operand type: expected *mut Coroutine(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), [i8, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), []), ()]), found *mut Coroutine(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), [i32, std::future::ResumeTy, (), (), CoroutineWitness(DefId(0:8 ~ async_fn_once_for_async_fn[6cdf]::main::{closure#0}::{closure#1}), []), ()]).See <https://github.com/rust-lang/rust/issues/114858>.

tests/ui/async-await/async-closures/def-path.rs

@@ -8,7 +8,7 @@ fn main() {
     //~^ NOTE the expected `async` closure body
     let () = x();
     //~^ ERROR mismatched types
-    //~| NOTE this expression has type `{static main::{closure#0}::{closure#0} upvar_tys=
+    //~| NOTE this expression has type `{static main::{closure#0}::{closure#0}<
     //~| NOTE expected `async` closure body, found `()`
-    //~| NOTE expected `async` closure body `{static main::{closure#0}::{closure#0}
+    //~| NOTE expected `async` closure body `{static main::{closure#0}::{closure#0}<
 }

tests/ui/async-await/async-closures/def-path.stderr

@@ -5,11 +5,11 @@ LL |     let x = async || {};
    |                      -- the expected `async` closure body
 LL |
 LL |     let () = x();
-   |         ^^   --- this expression has type `{static main::{closure#0}::{closure#0} upvar_tys=?7t witness=?8t}`
+   |         ^^   --- this expression has type `{static main::{closure#0}::{closure#0}<?7t> upvar_tys=?15t witness=?6t}`
    |         |
    |         expected `async` closure body, found `()`
    |
-   = note: expected `async` closure body `{static main::{closure#0}::{closure#0} upvar_tys=?7t witness=?8t}`
+   = note: expected `async` closure body `{static main::{closure#0}::{closure#0}<?7t> upvar_tys=?15t witness=?6t}`
                          found unit type `()`
 
 error: aborting due to 1 previous error

tests/ui/async-await/async-closures/drop.rs

@@ -0,0 +1,40 @@
+// aux-build:block-on.rs
+// edition:2018
+// run-pass
+// check-run-results
+
+#![feature(async_closure, async_fn_traits)]
+#![allow(unused)]
+
+extern crate block_on;
+
+use std::ops::AsyncFnOnce;
+
+struct DropMe(i32);
+
+impl Drop for DropMe {
+    fn drop(&mut self) {
+        println!("{} was dropped", self.0);
+    }
+}
+
+async fn call_once(f: impl AsyncFnOnce()) {
+    println!("before call");
+    let fut = Box::pin(f());
+    println!("after call");
+    drop(fut);
+    println!("future dropped");
+}
+
+fn main() {
+    block_on::block_on(async {
+        let d = DropMe(42);
+        let async_closure = async move || {
+            let d = &d;
+            println!("called");
+        };
+
+        call_once(async_closure).await;
+        println!("after");
+    });
+}

tests/ui/async-await/async-closures/drop.run.stdout

@@ -0,0 +1,5 @@
+before call
+after call
+42 was dropped
+future dropped
+after

tests/ui/async-await/async-closures/mangle.rs

@@ -0,0 +1,36 @@
+// aux-build:block-on.rs
+// edition:2021
+// build-pass
+// revisions: v0 legacy
+//[v0] compile-flags: -Csymbol-mangling-version=v0
+//[legacy] compile-flags: -Csymbol-mangling-version=legacy -Zunstable-options
+
+// FIXME(async_closures): When `fn_sig_for_fn_abi` is fixed, remove this.
+// ignore-pass (test emits codegen-time warnings)
+
+#![feature(async_closure, noop_waker, async_fn_traits)]
+
+extern crate block_on;
+
+use std::future::Future;
+use std::ops::{AsyncFnMut, AsyncFnOnce};
+use std::pin::pin;
+use std::task::*;
+
+async fn call_mut(f: &mut impl AsyncFnMut()) {
+    f().await;
+}
+
+async fn call_once(f: impl AsyncFnOnce()) {
+    f().await;
+}
+
+fn main() {
+    block_on::block_on(async {
+        let mut async_closure = async move || {
+            println!("called");
+        };
+        call_mut(&mut async_closure).await;
+        call_once(async_closure).await;
+    });
+}