From f1fef64e19909487ff2640bce58ce49fcfb4b85d Mon Sep 17 00:00:00 2001 From: Michael Goulet Date: Wed, 28 Feb 2024 20:25:25 +0000 Subject: Fix ABI for FnMut/Fn impls for async closures --- compiler/rustc_middle/src/mir/visit.rs | 1 + compiler/rustc_middle/src/ty/instance.rs | 11 ++++++++++- compiler/rustc_mir_transform/src/shim.rs | 24 +++++++++++++++++++----- compiler/rustc_ty_utils/src/abi.rs | 15 +++++++++++---- compiler/rustc_ty_utils/src/instance.rs | 2 ++ 5 files changed, 43 insertions(+), 10 deletions(-) (limited to 'compiler') diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs index 562aed5a643..be960669ff4 100644 --- a/compiler/rustc_middle/src/mir/visit.rs +++ b/compiler/rustc_middle/src/mir/visit.rs @@ -347,6 +347,7 @@ macro_rules! make_mir_visitor { ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } | 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) => {} diff --git a/compiler/rustc_middle/src/ty/instance.rs b/compiler/rustc_middle/src/ty/instance.rs index bbe0915baa2..18ef4ed549b 100644 --- a/compiler/rustc_middle/src/ty/instance.rs +++ b/compiler/rustc_middle/src/ty/instance.rs @@ -95,7 +95,15 @@ pub enum InstanceDef<'tcx> { /// 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 }, + ConstructCoroutineInClosureShim { + coroutine_closure_def_id: DefId, + // 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) -> ::Output`, that is to say + // that it returns the `FnOnce`-flavored coroutine but takes the closure + // by ref (and similarly for `Fn::call`). + receiver_by_ref: bool, + }, /// `<[coroutine] as Future>::poll`, but for coroutines produced when `AsyncFnOnce` /// is called on a coroutine-closure whose closure kind greater than `FnOnce`, or @@ -188,6 +196,7 @@ impl<'tcx> InstanceDef<'tcx> { | InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ } | ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id: def_id, + receiver_by_ref: _, } | ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id } | InstanceDef::DropGlue(def_id, _) diff --git a/compiler/rustc_mir_transform/src/shim.rs b/compiler/rustc_mir_transform/src/shim.rs index 3efaa69a7e7..4b2243598dc 100644 --- a/compiler/rustc_mir_transform/src/shim.rs +++ b/compiler/rustc_mir_transform/src/shim.rs @@ -70,9 +70,10 @@ 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 } => { - build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id) - } + ty::InstanceDef::ConstructCoroutineInClosureShim { + coroutine_closure_def_id, + receiver_by_ref, + } => build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id, receiver_by_ref), ty::InstanceDef::CoroutineKindShim { coroutine_def_id } => { return tcx.optimized_mir(coroutine_def_id).coroutine_by_move_body().unwrap().clone(); @@ -1015,12 +1016,17 @@ fn build_fn_ptr_addr_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'t fn build_construct_coroutine_by_move_shim<'tcx>( tcx: TyCtxt<'tcx>, coroutine_closure_def_id: DefId, + receiver_by_ref: bool, ) -> Body<'tcx> { - let self_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity(); + let mut self_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity(); let ty::CoroutineClosure(_, args) = *self_ty.kind() else { bug!(); }; + if receiver_by_ref { + self_ty = Ty::new_mut_ptr(tcx, self_ty); + } + 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()), @@ -1076,11 +1082,19 @@ fn build_construct_coroutine_by_move_shim<'tcx>( let source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id, + receiver_by_ref, }); 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(())); + dump_mir( + tcx, + false, + if receiver_by_ref { "coroutine_closure_by_ref" } else { "coroutine_closure_by_move" }, + &0, + &body, + |_, _| Ok(()), + ); body } diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs index 7d54083fbd5..baf4de768c5 100644 --- a/compiler/rustc_ty_utils/src/abi.rs +++ b/compiler/rustc_ty_utils/src/abi.rs @@ -118,11 +118,18 @@ fn fn_sig_for_fn_abi<'tcx>( // a separate def-id for these bodies. let mut coroutine_kind = args.as_coroutine_closure().kind(); - if let InstanceDef::ConstructCoroutineInClosureShim { .. } = instance.def { - coroutine_kind = ty::ClosureKind::FnOnce; - } + let env_ty = + if let InstanceDef::ConstructCoroutineInClosureShim { receiver_by_ref, .. } = + instance.def + { + coroutine_kind = ty::ClosureKind::FnOnce; - let env_ty = tcx.closure_env_ty(coroutine_ty, coroutine_kind, env_region); + // Implementations of `FnMut` and `Fn` for coroutine-closures + // still take their receiver by ref. + if receiver_by_ref { Ty::new_mut_ptr(tcx, coroutine_ty) } else { coroutine_ty } + } else { + tcx.closure_env_ty(coroutine_ty, coroutine_kind, env_region) + }; let sig = sig.skip_binder(); ty::Binder::bind_with_vars( diff --git a/compiler/rustc_ty_utils/src/instance.rs b/compiler/rustc_ty_utils/src/instance.rs index c2ea89f4c29..a8f9afb87dd 100644 --- a/compiler/rustc_ty_utils/src/instance.rs +++ b/compiler/rustc_ty_utils/src/instance.rs @@ -282,6 +282,7 @@ fn resolve_associated_item<'tcx>( Some(Instance { def: ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id, + receiver_by_ref: target_kind != ty::ClosureKind::FnOnce, }, args, }) @@ -304,6 +305,7 @@ fn resolve_associated_item<'tcx>( Some(Instance { def: ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id, + receiver_by_ref: false, }, args, }) -- cgit 1.4.1-3-g733a5