Auto merge of #120361 - compiler-errors:async-closures, r=oli-obk

Rework support for async closures; allow them to return futures that borrow from the closure's captures

This PR implements a new lowering for async closures via `TyKind::CoroutineClosure` which handles the curious relationship between the closure and the coroutine that it returns.

I wrote up a bunch in [this hackmd](https://hackmd.io/`@compiler-errors/S1HvqQxca)` which will be copied to the dev guide after this PR lands, and hopefully left sufficient comments in the source code explaining why this change is as large as it is.

This also necessitates that they begin implementing the `AsyncFn`-family of traits, rather than the `Fn`-family of traits -- if you need `Fn` implementations, you should probably use the non-sugar `|| async {}` syntax instead.

Notably this PR does not yet implement `async Fn()` syntax sugar for bounds, but I expect to add those soon (**edit:** #120392). For now, users must use `AsyncFn()` traits directly, which necessitates adding the `async_fn_traits` feature gate as well. I will add this as a follow-up very soon.

r? oli-obk

This is based on top of #120322, but that PR is minimal.

15 files changed, 373 insertions, 9 deletions

diff --git a/compiler/rustc_mir_transform/src/abort_unwinding_calls.rs b/compiler/rustc_mir_transform/src/abort_unwinding_calls.rs
index dfc7a9891f9..451d3be255f 100644
--- a/compiler/rustc_mir_transform/src/abort_unwinding_calls.rs
+++ b/compiler/rustc_mir_transform/src/abort_unwinding_calls.rs
@@ -39,6 +39,7 @@ impl<'tcx> MirPass<'tcx> for AbortUnwindingCalls {
         let body_abi = match body_ty.kind() {
             ty::FnDef(..) => body_ty.fn_sig(tcx).abi(),
             ty::Closure(..) => Abi::RustCall,
+            ty::CoroutineClosure(..) => Abi::RustCall,
             ty::Coroutine(..) => Abi::Rust,
             _ => span_bug!(body.span, "unexpected body ty: {:?}", body_ty),
         };
diff --git a/compiler/rustc_mir_transform/src/check_unsafety.rs b/compiler/rustc_mir_transform/src/check_unsafety.rs
index 6c4c3917cb5..fbb62695383 100644
--- a/compiler/rustc_mir_transform/src/check_unsafety.rs
+++ b/compiler/rustc_mir_transform/src/check_unsafety.rs
@@ -128,7 +128,9 @@ impl<'tcx> Visitor<'tcx> for UnsafetyChecker<'_, 'tcx> {
                         ),
                     }
                 }
-                &AggregateKind::Closure(def_id, _) | &AggregateKind::Coroutine(def_id, _) => {
+                &AggregateKind::Closure(def_id, _)
+                | &AggregateKind::CoroutineClosure(def_id, _)
+                | &AggregateKind::Coroutine(def_id, _) => {
                     let def_id = def_id.expect_local();
                     let UnsafetyCheckResult { violations, used_unsafe_blocks, .. } =
                         self.tcx.mir_unsafety_check_result(def_id);
diff --git a/compiler/rustc_mir_transform/src/const_prop_lint.rs b/compiler/rustc_mir_transform/src/const_prop_lint.rs
index aa22b8c7c58..f2448ee3d44 100644
--- a/compiler/rustc_mir_transform/src/const_prop_lint.rs
+++ b/compiler/rustc_mir_transform/src/const_prop_lint.rs
@@ -607,7 +607,8 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
                     AggregateKind::Array(_)
                     | AggregateKind::Tuple
                     | AggregateKind::Closure(_, _)
-                    | AggregateKind::Coroutine(_, _) => VariantIdx::new(0),
+                    | AggregateKind::Coroutine(_, _)
+                    | AggregateKind::CoroutineClosure(_, _) => VariantIdx::new(0),
                 },
             },
 
diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index bde879f6067..297b2fa143d 100644
--- a/compiler/rustc_mir_transform/src/coroutine.rs
+++ b/compiler/rustc_mir_transform/src/coroutine.rs
@@ -50,6 +50,9 @@
 //! For coroutines with state 1 (returned) and state 2 (poisoned) it does nothing.
 //! Otherwise it drops all the values in scope at the last suspension point.
 
+mod by_move_body;
+pub use by_move_body::ByMoveBody;
+
 use crate::abort_unwinding_calls;
 use crate::deref_separator::deref_finder;
 use crate::errors;
diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
new file mode 100644
index 00000000000..fcd4715b9e8
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
@@ -0,0 +1,156 @@
+//! A MIR pass which duplicates a coroutine's body and removes any derefs which
+//! would be present for upvars that are taken by-ref. The result of which will
+//! be a coroutine body that takes all of its upvars by-move, and which we stash
+//! into the `CoroutineInfo` for all coroutines returned by coroutine-closures.
+
+use rustc_data_structures::fx::FxIndexSet;
+use rustc_hir as hir;
+use rustc_middle::mir::visit::MutVisitor;
+use rustc_middle::mir::{self, dump_mir, MirPass};
+use rustc_middle::ty::{self, InstanceDef, Ty, TyCtxt};
+use rustc_target::abi::FieldIdx;
+
+pub struct ByMoveBody;
+
+impl<'tcx> MirPass<'tcx> for ByMoveBody {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
+        let Some(coroutine_def_id) = body.source.def_id().as_local() else {
+            return;
+        };
+        let Some(hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure)) =
+            tcx.coroutine_kind(coroutine_def_id)
+        else {
+            return;
+        };
+        let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
+        let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!() };
+
+        let coroutine_kind = args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap();
+        if coroutine_kind == ty::ClosureKind::FnOnce {
+            return;
+        }
+
+        let mut by_ref_fields = FxIndexSet::default();
+        let by_move_upvars = Ty::new_tup_from_iter(
+            tcx,
+            tcx.closure_captures(coroutine_def_id).iter().enumerate().map(|(idx, capture)| {
+                if capture.is_by_ref() {
+                    by_ref_fields.insert(FieldIdx::from_usize(idx));
+                }
+                capture.place.ty()
+            }),
+        );
+        let by_move_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::FnOnce),
+                    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: by_move_upvars,
+                },
+            )
+            .args,
+        );
+
+        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::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);
+        }
+    }
+}
+
+struct MakeByMoveBody<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    by_ref_fields: FxIndexSet<FieldIdx>,
+    by_move_coroutine_ty: Ty<'tcx>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for MakeByMoveBody<'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn visit_place(
+        &mut self,
+        place: &mut mir::Place<'tcx>,
+        context: mir::visit::PlaceContext,
+        location: mir::Location,
+    ) {
+        if place.local == ty::CAPTURE_STRUCT_LOCAL
+            && !place.projection.is_empty()
+            && let mir::ProjectionElem::Field(idx, ty) = place.projection[0]
+            && self.by_ref_fields.contains(&idx)
+        {
+            let (begin, end) = place.projection[1..].split_first().unwrap();
+            // FIXME(async_closures): I'm actually a bit surprised to see that we always
+            // initially deref the by-ref upvars. If this is not actually true, then we
+            // will at least get an ICE that explains why this isn't true :^)
+            assert_eq!(*begin, mir::ProjectionElem::Deref);
+            // Peel one ref off of the ty.
+            let peeled_ty = ty.builtin_deref(true).unwrap().ty;
+            *place = mir::Place {
+                local: place.local,
+                projection: self.tcx.mk_place_elems_from_iter(
+                    [mir::ProjectionElem::Field(idx, peeled_ty)]
+                        .into_iter()
+                        .chain(end.iter().copied()),
+                ),
+            };
+        }
+        self.super_place(place, context, location);
+    }
+
+    fn visit_local_decl(&mut self, local: mir::Local, local_decl: &mut mir::LocalDecl<'tcx>) {
+        // Replace the type of the self arg.
+        if local == ty::CAPTURE_STRUCT_LOCAL {
+            local_decl.ty = self.by_move_coroutine_ty;
+        }
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
index 5b4d58836b4..01fae7c0bec 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -156,7 +156,9 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
 fn is_closure_or_coroutine(statement: &Statement<'_>) -> bool {
     match statement.kind {
         StatementKind::Assign(box (_, Rvalue::Aggregate(box ref agg_kind, _))) => match agg_kind {
-            AggregateKind::Closure(_, _) | AggregateKind::Coroutine(_, _) => true,
+            AggregateKind::Closure(_, _)
+            | AggregateKind::Coroutine(_, _)
+            | AggregateKind::CoroutineClosure(..) => true,
             _ => false,
         },
         _ => false,
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
index ad12bce9b02..6a37047a693 100644
--- a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
+++ b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -696,6 +696,7 @@ fn try_write_constant<'tcx>(
         | ty::Bound(..)
         | ty::Placeholder(..)
         | ty::Closure(..)
+        | ty::CoroutineClosure(..)
         | ty::Coroutine(..)
         | ty::Dynamic(..) => throw_machine_stop_str!("unsupported type"),
 
diff --git a/compiler/rustc_mir_transform/src/ffi_unwind_calls.rs b/compiler/rustc_mir_transform/src/ffi_unwind_calls.rs
index db7dfc5b43e..b0d758bcbfe 100644
--- a/compiler/rustc_mir_transform/src/ffi_unwind_calls.rs
+++ b/compiler/rustc_mir_transform/src/ffi_unwind_calls.rs
@@ -57,6 +57,7 @@ fn has_ffi_unwind_calls(tcx: TyCtxt<'_>, local_def_id: LocalDefId) -> bool {
     let body_abi = match body_ty.kind() {
         ty::FnDef(..) => body_ty.fn_sig(tcx).abi(),
         ty::Closure(..) => Abi::RustCall,
+        ty::CoroutineClosure(..) => Abi::RustCall,
         ty::Coroutine(..) => Abi::Rust,
         _ => span_bug!(body.span, "unexpected body ty: {:?}", body_ty),
     };
diff --git a/compiler/rustc_mir_transform/src/gvn.rs b/compiler/rustc_mir_transform/src/gvn.rs
index 390ec3e1a36..f9798bc4e70 100644
--- a/compiler/rustc_mir_transform/src/gvn.rs
+++ b/compiler/rustc_mir_transform/src/gvn.rs
@@ -861,9 +861,10 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
         let tcx = self.tcx;
         if fields.is_empty() {
             let is_zst = match *kind {
-                AggregateKind::Array(..) | AggregateKind::Tuple | AggregateKind::Closure(..) => {
-                    true
-                }
+                AggregateKind::Array(..)
+                | AggregateKind::Tuple
+                | AggregateKind::Closure(..)
+                | AggregateKind::CoroutineClosure(..) => true,
                 // Only enums can be non-ZST.
                 AggregateKind::Adt(did, ..) => tcx.def_kind(did) != DefKind::Enum,
                 // Coroutines are never ZST, as they at least contain the implicit states.
@@ -885,7 +886,9 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
                 assert!(!fields.is_empty());
                 (AggregateTy::Tuple, FIRST_VARIANT)
             }
-            AggregateKind::Closure(did, substs) | AggregateKind::Coroutine(did, substs) => {
+            AggregateKind::Closure(did, substs)
+            | AggregateKind::CoroutineClosure(did, substs)
+            | AggregateKind::Coroutine(did, substs) => {
                 (AggregateTy::Def(did, substs), FIRST_VARIANT)
             }
             AggregateKind::Adt(did, variant_index, substs, _, None) => {
diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs
index 67668a216de..e77553a03d6 100644
--- a/compiler/rustc_mir_transform/src/inline.rs
+++ b/compiler/rustc_mir_transform/src/inline.rs
@@ -317,6 +317,8 @@ impl<'tcx> Inliner<'tcx> {
             | InstanceDef::ReifyShim(_)
             | InstanceDef::FnPtrShim(..)
             | InstanceDef::ClosureOnceShim { .. }
+            | InstanceDef::ConstructCoroutineInClosureShim { .. }
+            | InstanceDef::CoroutineKindShim { .. }
             | InstanceDef::DropGlue(..)
             | InstanceDef::CloneShim(..)
             | InstanceDef::ThreadLocalShim(..)
diff --git a/compiler/rustc_mir_transform/src/inline/cycle.rs b/compiler/rustc_mir_transform/src/inline/cycle.rs
index d30e0bad813..5b03bc361dd 100644
--- a/compiler/rustc_mir_transform/src/inline/cycle.rs
+++ b/compiler/rustc_mir_transform/src/inline/cycle.rs
@@ -87,6 +87,8 @@ pub(crate) fn mir_callgraph_reachable<'tcx>(
                 | InstanceDef::ReifyShim(_)
                 | InstanceDef::FnPtrShim(..)
                 | InstanceDef::ClosureOnceShim { .. }
+                | InstanceDef::ConstructCoroutineInClosureShim { .. }
+                | InstanceDef::CoroutineKindShim { .. }
                 | InstanceDef::ThreadLocalShim { .. }
                 | InstanceDef::CloneShim(..) => {}
 
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 69f93fa3a0e..031515ea958 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -307,6 +307,10 @@ fn mir_const(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
             &Lint(check_packed_ref::CheckPackedRef),
             &Lint(check_const_item_mutation::CheckConstItemMutation),
             &Lint(function_item_references::FunctionItemReferences),
+            // If this is an async closure's output coroutine, generate
+            // by-move and by-mut bodies if needed. We do this first so
+            // they can be optimized in lockstep with their parent bodies.
+            &coroutine::ByMoveBody,
             // What we need to do constant evaluation.
             &simplify::SimplifyCfg::Initial,
             &rustc_peek::SanityCheck, // Just a lint
diff --git a/compiler/rustc_mir_transform/src/pass_manager.rs b/compiler/rustc_mir_transform/src/pass_manager.rs
index c1ef2b9f887..605e1ad46d7 100644
--- a/compiler/rustc_mir_transform/src/pass_manager.rs
+++ b/compiler/rustc_mir_transform/src/pass_manager.rs
@@ -189,6 +189,15 @@ fn run_passes_inner<'tcx>(
 
         body.pass_count = 1;
     }
+
+    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);
+        }
+    }
 }
 
 pub fn validate_body<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, when: String) {
diff --git a/compiler/rustc_mir_transform/src/remove_zsts.rs b/compiler/rustc_mir_transform/src/remove_zsts.rs
index 34d57a45301..9a94cae3382 100644
--- a/compiler/rustc_mir_transform/src/remove_zsts.rs
+++ b/compiler/rustc_mir_transform/src/remove_zsts.rs
@@ -46,6 +46,7 @@ fn maybe_zst(ty: Ty<'_>) -> bool {
         ty::Adt(..)
         | ty::Array(..)
         | ty::Closure(..)
+        | ty::CoroutineClosure(..)
         | ty::Tuple(..)
         | ty::Alias(ty::Opaque, ..) => true,
         // definitely ZST
diff --git a/compiler/rustc_mir_transform/src/shim.rs b/compiler/rustc_mir_transform/src/shim.rs
index 89414ce940e..7b6de3a5439 100644
--- a/compiler/rustc_mir_transform/src/shim.rs
+++ b/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,11 +66,76 @@ 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::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::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);
 
@@ -981,3 +1046,114 @@ fn build_fn_ptr_addr_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'t
     let source = MirSource::from_instance(ty::InstanceDef::FnPtrAddrShim(def_id, self_ty));
     new_body(source, IndexVec::from_elem_n(start_block, 1), locals, sig.inputs().len(), span)
 }
+
+fn build_construct_coroutine_by_move_shim<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    coroutine_closure_def_id: DefId,
+) -> Body<'tcx> {
+    let self_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity();
+    let ty::CoroutineClosure(_, args) = *self_ty.kind() else {
+        bug!();
+    };
+
+    let poly_sig = args.as_coroutine_closure().coroutine_closure_sig().map_bound(|sig| {
+        tcx.mk_fn_sig(
+            [self_ty].into_iter().chain(sig.tupled_inputs_ty.tuple_fields()),
+            sig.to_coroutine_given_kind_and_upvars(
+                tcx,
+                args.as_coroutine_closure().parent_args(),
+                tcx.coroutine_for_closure(coroutine_closure_def_id),
+                ty::ClosureKind::FnOnce,
+                tcx.lifetimes.re_erased,
+                args.as_coroutine_closure().tupled_upvars_ty(),
+                args.as_coroutine_closure().coroutine_captures_by_ref_ty(),
+            ),
+            sig.c_variadic,
+            sig.unsafety,
+            sig.abi,
+        )
+    });
+    let sig = tcx.liberate_late_bound_regions(coroutine_closure_def_id, poly_sig);
+    let ty::Coroutine(coroutine_def_id, coroutine_args) = *sig.output().kind() else {
+        bug!();
+    };
+
+    let span = tcx.def_span(coroutine_closure_def_id);
+    let locals = local_decls_for_sig(&sig, span);
+
+    let mut fields = vec![];
+    for idx in 1..sig.inputs().len() {
+        fields.push(Operand::Move(Local::from_usize(idx + 1).into()));
+    }
+    for (idx, ty) in args.as_coroutine_closure().upvar_tys().iter().enumerate() {
+        fields.push(Operand::Move(tcx.mk_place_field(
+            Local::from_usize(1).into(),
+            FieldIdx::from_usize(idx),
+            ty,
+        )));
+    }
+
+    let source_info = SourceInfo::outermost(span);
+    let rvalue = Rvalue::Aggregate(
+        Box::new(AggregateKind::Coroutine(coroutine_def_id, coroutine_args)),
+        IndexVec::from_raw(fields),
+    );
+    let stmt = Statement {
+        source_info,
+        kind: StatementKind::Assign(Box::new((Place::return_place(), rvalue))),
+    };
+    let statements = vec![stmt];
+    let start_block = BasicBlockData {
+        statements,
+        terminator: Some(Terminator { source_info, kind: TerminatorKind::Return }),
+        is_cleanup: false,
+    };
+
+    let source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim {
+        coroutine_closure_def_id,
+        target_kind: ty::ClosureKind::FnOnce,
+    });
+
+    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>(
+    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
+}