Cap closure kind, get rid of by_mut_body

Author: compiler-errors

compiler/rustc_borrowck/src/type_check/input_output.rs

@@ -88,7 +88,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                     self.tcx(),
                     ty::CoroutineArgsParts {
                         parent_args: args.parent_args(),
-                        kind_ty: Ty::from_closure_kind(self.tcx(), args.kind()),
+                        kind_ty: Ty::from_closure_kind(self.tcx(), args.kind().cap_for_coroutine()),
                         return_ty: user_provided_sig.output(),
                         tupled_upvars_ty,
                         // For async closures, none of these can be annotated, so just fill

compiler/rustc_borrowck/src/universal_regions.rs

@@ -755,7 +755,7 @@ impl<'cx, 'tcx> UniversalRegionsBuilder<'cx, 'tcx> {
                     tcx,
                     args.as_coroutine_closure().parent_args(),
                     tcx.coroutine_for_closure(def_id),
-                    closure_kind,
+                    closure_kind.cap_for_coroutine(),
                     env_region,
                     args.as_coroutine_closure().tupled_upvars_ty(),
                     args.as_coroutine_closure().coroutine_captures_by_ref_ty(),

compiler/rustc_hir_typeck/src/callee.rs

@@ -183,16 +183,20 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     kind: TypeVariableOriginKind::TypeInference,
                     span: callee_expr.span,
                 });
+                // We may actually receive a coroutine back whose kind is different
+                // from the closure that this dispatched from. This is because when
+                // we have no captures, we automatically implement `FnOnce`. This
+                // impl forces the closure kind to `FnOnce` i.e. `u8`.
+                let kind_ty = self.next_ty_var(TypeVariableOrigin {
+                    kind: TypeVariableOriginKind::TypeInference,
+                    span: callee_expr.span,
+                });
                 let call_sig = self.tcx.mk_fn_sig(
                     [coroutine_closure_sig.tupled_inputs_ty],
                     coroutine_closure_sig.to_coroutine(
                         self.tcx,
                         closure_args.parent_args(),
-                        // Inherit the kind ty of the closure, since we're calling this
-                        // coroutine with the most relaxed `AsyncFn*` trait that we can.
-                        // We don't necessarily need to do this here, but it saves us
-                        // computing one more infer var that will get constrained later.
-                        closure_args.kind_ty(),
+                        kind_ty,
                         self.tcx.coroutine_for_closure(def_id),
                         tupled_upvars_ty,
                     ),

compiler/rustc_hir_typeck/src/closure.rs

@@ -261,6 +261,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     },
                 );
 
+                let coroutine_kind_ty = self.next_ty_var(TypeVariableOrigin {
+                    // FIXME(eddyb) distinguish closure kind inference variables from the rest.
+                    kind: TypeVariableOriginKind::ClosureSynthetic,
+                    span: expr_span,
+                });
                 let coroutine_upvars_ty = self.next_ty_var(TypeVariableOrigin {
                     kind: TypeVariableOriginKind::ClosureSynthetic,
                     span: expr_span,
@@ -278,7 +283,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         sig.to_coroutine(
                             tcx,
                             parent_args,
-                            closure_kind_ty,
+                            coroutine_kind_ty,
                             tcx.coroutine_for_closure(expr_def_id),
                             coroutine_upvars_ty,
                         )

compiler/rustc_hir_typeck/src/upvar.rs

@@ -410,7 +410,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             self.demand_eqtype(
                 span,
                 coroutine_args.as_coroutine().kind_ty(),
-                Ty::from_closure_kind(self.tcx, closure_kind),
+                Ty::from_closure_kind(self.tcx, closure_kind.cap_for_coroutine()),
             );
         }
 

compiler/rustc_middle/src/mir/mod.rs

@@ -274,13 +274,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>>,
 
@@ -301,7 +294,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,
         }
@@ -616,10 +608,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)

compiler/rustc_middle/src/mir/visit.rs

@@ -345,8 +345,8 @@ 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 } |
+                        ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id } |
                         ty::InstanceDef::DropGlue(_def_id, None) => {}
 
                         ty::InstanceDef::FnPtrShim(_def_id, ty) |

compiler/rustc_middle/src/ty/instance.rs

@@ -90,24 +90,20 @@ 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,
-    },
+    ConstructCoroutineInClosureShim { coroutine_closure_def_id: DefId },
 
     /// `<[coroutine] as Future>::poll`, but for coroutines produced when `AsyncFnOnce`
     /// 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.
-    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 +188,8 @@ impl<'tcx> InstanceDef<'tcx> {
             | 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::CoroutineKindShim { coroutine_def_id: def_id }
             | InstanceDef::DropGlue(def_id, _)
             | InstanceDef::CloneShim(def_id, _)
             | InstanceDef::FnPtrAddrShim(def_id, _) => def_id,
@@ -654,10 +649,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,
                 })
             }

compiler/rustc_middle/src/ty/sty.rs

@@ -485,7 +485,7 @@ impl<'tcx> CoroutineClosureSignature<'tcx> {
         self.to_coroutine(
             tcx,
             parent_args,
-            Ty::from_closure_kind(tcx, goal_kind),
+            Ty::from_closure_kind(tcx, goal_kind.cap_for_coroutine()),
             coroutine_def_id,
             tupled_upvars_ty,
         )

compiler/rustc_mir_transform/src/coroutine/by_move_body.rs

@@ -64,45 +64,10 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
         by_move_body.source = mir::MirSource {
             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/pass_manager.rs

@@ -194,9 +194,6 @@ fn run_passes_inner<'tcx>(
         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);
-        }
     }
 }
 

compiler/rustc_mir_transform/src/shim.rs

@@ -3,8 +3,8 @@ use rustc_hir::def_id::DefId;
 use rustc_hir::lang_items::LangItem;
 use rustc_middle::mir::*;
 use rustc_middle::query::Providers;
+use rustc_middle::ty::GenericArgs;
 use rustc_middle::ty::{self, CoroutineArgs, EarlyBinder, Ty, TyCtxt};
-use rustc_middle::ty::{GenericArgs, CAPTURE_STRUCT_LOCAL};
 use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT};
 
 use rustc_index::{Idx, IndexVec};
@@ -66,39 +66,13 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
             build_call_shim(tcx, instance, Some(Adjustment::RefMut), CallKind::Direct(call_mut))
         }
 
-        ty::InstanceDef::ConstructCoroutineInClosureShim {
-            coroutine_closure_def_id,
-            target_kind,
-        } => match target_kind {
-            ty::ClosureKind::Fn => unreachable!("shouldn't be building shim for Fn"),
-            ty::ClosureKind::FnMut => {
-                // No need to optimize the body, it has already been optimized
-                // since we steal it from the `AsyncFn::call` body and just fix
-                // 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::ConstructCoroutineInClosureShim { coroutine_closure_def_id } => {
+            build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id)
+        }
 
-        ty::InstanceDef::CoroutineKindShim { coroutine_def_id, target_kind } => match target_kind {
-            ty::ClosureKind::Fn => unreachable!(),
-            ty::ClosureKind::FnMut => {
-                return tcx
-                    .optimized_mir(coroutine_def_id)
-                    .coroutine_by_mut_body()
-                    .unwrap()
-                    .clone();
-            }
-            ty::ClosureKind::FnOnce => {
-                return tcx
-                    .optimized_mir(coroutine_def_id)
-                    .coroutine_by_move_body()
-                    .unwrap()
-                    .clone();
-            }
-        },
+        ty::InstanceDef::CoroutineKindShim { coroutine_def_id } => {
+            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
@@ -119,21 +93,11 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
                 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(),
-                    }
+                    assert_eq!(
+                        args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
+                        ty::ClosureKind::FnOnce
+                    );
+                    coroutine_body.coroutine_by_move_body().unwrap().coroutine_drop().unwrap()
                 };
 
                 let mut body = EarlyBinder::bind(body.clone()).instantiate(tcx, args);
@@ -1111,7 +1075,6 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
 
     let source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim {
         coroutine_closure_def_id,
-        target_kind: ty::ClosureKind::FnOnce,
     });
 
     let body =
@@ -1120,40 +1083,3 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
 
     body
 }
-
-fn build_construct_coroutine_by_mut_shim<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    coroutine_closure_def_id: DefId,
-) -> Body<'tcx> {
-    let mut body = tcx.optimized_mir(coroutine_closure_def_id).clone();
-    let coroutine_closure_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity();
-    let ty::CoroutineClosure(_, args) = *coroutine_closure_ty.kind() else {
-        bug!();
-    };
-    let args = args.as_coroutine_closure();
-
-    body.local_decls[RETURN_PLACE].ty =
-        tcx.instantiate_bound_regions_with_erased(args.coroutine_closure_sig().map_bound(|sig| {
-            sig.to_coroutine_given_kind_and_upvars(
-                tcx,
-                args.parent_args(),
-                tcx.coroutine_for_closure(coroutine_closure_def_id),
-                ty::ClosureKind::FnMut,
-                tcx.lifetimes.re_erased,
-                args.tupled_upvars_ty(),
-                args.coroutine_captures_by_ref_ty(),
-            )
-        }));
-    body.local_decls[CAPTURE_STRUCT_LOCAL].ty =
-        Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, coroutine_closure_ty);
-
-    body.source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim {
-        coroutine_closure_def_id,
-        target_kind: ty::ClosureKind::FnMut,
-    });
-
-    body.pass_count = 0;
-    dump_mir(tcx, false, "coroutine_closure_by_mut", &0, &body, |_, _| Ok(()));
-
-    body
-}

compiler/rustc_symbol_mangling/src/legacy.rs

@@ -76,16 +76,10 @@ pub(super) fn mangle<'tcx>(
         }
         // 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();
-            }
-        },
+        ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
+        | ty::InstanceDef::CoroutineKindShim { .. } => {
+            printer.write_str("{{fn-once-shim}}").unwrap();
+        }
         _ => {}
     }