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

Author: compiler-errors

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
-    /// will select the body that is produced by the `ByMoveBody` transform, and thus
+    /// is called on a coroutine-closure whose closure kind greater than `FnOnce`, or
+    /// 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()
-        && let Some(by_move_body) = coroutine.by_move_body.as_mut()
-    {
-        run_passes_inner(tcx, by_move_body, passes, phase_change, validate_each);
+    if let Some(coroutine) = body.coroutine.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);
+        }
+        if let Some(by_mut_body) = coroutine.by_mut_body.as_mut() {
+            run_passes_inner(tcx, by_mut_body, passes, phase_change, validate_each);
+        }
     }
 }