Auto merge of #120717 - compiler-errors:cap-closure-kind, r=oli-obk

For async closures, cap closure kind, get rid of `by_mut_body`

Right now we have three `AsyncFn*` traits, and three corresponding futures that are returned by the `call_*` functions for them. This is fine, but it is a bit excessive, since the future returned by `AsyncFn` and `AsyncFnMut` are identical. Really, the only distinction we need to make with these bodies is "by ref" and "by move".

This PR removes `AsyncFn::CallFuture` and renames `AsyncFnMut::CallMutFuture` to `AsyncFnMut::CallRefFuture`. This simplifies MIR building for async closures, since we don't need to build an extra "by mut" body, but just a "by move" body which is materially different.

We need to do a bit of delicate handling of the ClosureKind for async closures, since we need to "cap" it to `AsyncFnMut` in some cases when we only care about what body we're looking for.

This also fixes a bug where `<{async closure} as Fn>::call` was returning a body that takes the async-closure receiver *by move*.

This also helps align the `AsyncFn` traits to the `LendingFn` traits' eventual designs.

4 files changed, 36 insertions, 25 deletions

diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index d57ffc0f8b5..984d4687ef8 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -278,13 +278,6 @@ 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>>,
 
@@ -305,7 +298,6 @@ 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,
         }
@@ -628,10 +620,6 @@ impl<'tcx> 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)
diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs
index 845b1717550..be960669ff4 100644
--- a/compiler/rustc_middle/src/mir/visit.rs
+++ b/compiler/rustc_middle/src/mir/visit.rs
@@ -345,8 +345,11 @@ macro_rules! make_mir_visitor {
                         ty::InstanceDef::Virtual(_def_id, _) |
                         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::CoroutineKindShim { coroutine_def_id: _def_id, target_kind: _ } |
+                        ty::InstanceDef::ConstructCoroutineInClosureShim {
+                            coroutine_closure_def_id: _def_id,
+                            receiver_by_ref: _,
+                        } |
+                        ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id } |
                         ty::InstanceDef::DropGlue(_def_id, None) => {}
 
                         ty::InstanceDef::FnPtrShim(_def_id, ty) |
diff --git a/compiler/rustc_middle/src/ty/instance.rs b/compiler/rustc_middle/src/ty/instance.rs
index 814c3629b08..4748e961019 100644
--- a/compiler/rustc_middle/src/ty/instance.rs
+++ b/compiler/rustc_middle/src/ty/instance.rs
@@ -90,15 +90,19 @@ pub enum InstanceDef<'tcx> {
     /// and dispatch to the `FnMut::call_mut` instance for the closure.
     ClosureOnceShim { call_once: DefId, track_caller: bool },
 
-    /// `<[FnMut/Fn coroutine-closure] as FnOnce>::call_once` or
-    /// `<[Fn coroutine-closure] as FnMut>::call_mut`.
+    /// `<[FnMut/Fn coroutine-closure] as FnOnce>::call_once`
     ///
     /// The body generated here differs significantly from the `ClosureOnceShim`,
     /// since we need to generate a distinct coroutine type that will move the
     /// closure's upvars *out* of the closure.
     ConstructCoroutineInClosureShim {
         coroutine_closure_def_id: DefId,
-        target_kind: ty::ClosureKind,
+        // Whether the generated MIR body takes the coroutine by-ref. This is
+        // because the signature of `<{async fn} as FnMut>::call_mut` is:
+        // `fn(&mut self, args: A) -> <Self as FnOnce>::Output`, that is to say
+        // that it returns the `FnOnce`-flavored coroutine but takes the closure
+        // by mut ref (and similarly for `Fn::call`).
+        receiver_by_ref: bool,
     },
 
     /// `<[coroutine] as Future>::poll`, but for coroutines produced when `AsyncFnOnce`
@@ -107,7 +111,7 @@ pub enum InstanceDef<'tcx> {
     ///
     /// 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.
-    CoroutineKindShim { coroutine_def_id: DefId, target_kind: ty::ClosureKind },
+    CoroutineKindShim { coroutine_def_id: DefId },
 
     /// 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,9 +196,9 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ }
             | ty::InstanceDef::ConstructCoroutineInClosureShim {
                 coroutine_closure_def_id: def_id,
-                target_kind: _,
+                receiver_by_ref: _,
             }
-            | ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id, target_kind: _ }
+            | ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id }
             | InstanceDef::DropGlue(def_id, _)
             | InstanceDef::CloneShim(def_id, _)
             | InstanceDef::FnPtrAddrShim(def_id, _) => def_id,
@@ -651,10 +655,7 @@ impl<'tcx> Instance<'tcx> {
                 Some(Instance { def: ty::InstanceDef::Item(coroutine_def_id), args })
             } else {
                 Some(Instance {
-                    def: ty::InstanceDef::CoroutineKindShim {
-                        coroutine_def_id,
-                        target_kind: args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
-                    },
+                    def: ty::InstanceDef::CoroutineKindShim { coroutine_def_id },
                     args,
                 })
             }
diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs
index 540936d7d8a..6e0a9eb86dd 100644
--- a/compiler/rustc_middle/src/ty/sty.rs
+++ b/compiler/rustc_middle/src/ty/sty.rs
@@ -483,7 +483,7 @@ impl<'tcx> CoroutineClosureSignature<'tcx> {
         self.to_coroutine(
             tcx,
             parent_args,
-            Ty::from_closure_kind(tcx, goal_kind),
+            Ty::from_coroutine_closure_kind(tcx, goal_kind),
             coroutine_def_id,
             tupled_upvars_ty,
         )
@@ -2456,6 +2456,25 @@ impl<'tcx> Ty<'tcx> {
         }
     }
 
+    /// Like [`Ty::to_opt_closure_kind`], but it caps the "maximum" closure kind
+    /// to `FnMut`. This is because although we have three capability states,
+    /// `AsyncFn`/`AsyncFnMut`/`AsyncFnOnce`, we only need to distinguish two coroutine
+    /// bodies: by-ref and by-value.
+    ///
+    /// See the definition of `AsyncFn` and `AsyncFnMut` and the `CallRefFuture`
+    /// associated type for why we don't distinguish [`ty::ClosureKind::Fn`] and
+    /// [`ty::ClosureKind::FnMut`] for the purpose of the generated MIR bodies.
+    ///
+    /// This method should be used when constructing a `Coroutine` out of a
+    /// `CoroutineClosure`, when the `Coroutine`'s `kind` field is being populated
+    /// directly from the `CoroutineClosure`'s `kind`.
+    pub fn from_coroutine_closure_kind(tcx: TyCtxt<'tcx>, kind: ty::ClosureKind) -> Ty<'tcx> {
+        match kind {
+            ty::ClosureKind::Fn | ty::ClosureKind::FnMut => tcx.types.i16,
+            ty::ClosureKind::FnOnce => tcx.types.i32,
+        }
+    }
+
     /// Fast path helper for testing if a type is `Sized`.
     ///
     /// Returning true means the type is known to be sized. Returning