diff options
Diffstat (limited to 'compiler/rustc_mir_transform/src')
| -rw-r--r-- | compiler/rustc_mir_transform/src/abort_unwinding_calls.rs | 1 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/check_unsafety.rs | 4 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/const_prop_lint.rs | 3 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/coroutine.rs | 3 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/coroutine/by_move_body.rs | 156 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs | 4 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/dataflow_const_prop.rs | 1 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/ffi_unwind_calls.rs | 1 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/gvn.rs | 11 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/inline.rs | 2 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/inline/cycle.rs | 2 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/lib.rs | 4 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/pass_manager.rs | 9 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/remove_zsts.rs | 1 | ||||
| -rw-r--r-- | compiler/rustc_mir_transform/src/shim.rs | 180 |
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 +} |