From c567eddec2c628d4f13707866731e1b2013ad236 Mon Sep 17 00:00:00 2001 From: Michael Goulet Date: Wed, 24 Jan 2024 18:01:56 +0000 Subject: Add CoroutineClosure to TyKind, AggregateKind, UpvarArgs --- compiler/rustc_mir_dataflow/src/elaborate_drops.rs | 3 +++ compiler/rustc_mir_dataflow/src/move_paths/builder.rs | 8 ++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) (limited to 'compiler/rustc_mir_dataflow') diff --git a/compiler/rustc_mir_dataflow/src/elaborate_drops.rs b/compiler/rustc_mir_dataflow/src/elaborate_drops.rs index 862876f53c7..1b2f2cd9477 100644 --- a/compiler/rustc_mir_dataflow/src/elaborate_drops.rs +++ b/compiler/rustc_mir_dataflow/src/elaborate_drops.rs @@ -861,6 +861,9 @@ where let ty = self.place_ty(self.place); match ty.kind() { ty::Closure(_, args) => self.open_drop_for_tuple(args.as_closure().upvar_tys()), + ty::CoroutineClosure(_, args) => { + self.open_drop_for_tuple(args.as_coroutine_closure().upvar_tys()) + } // Note that `elaborate_drops` only drops the upvars of a coroutine, // and this is ok because `open_drop` here can only be reached // within that own coroutine's resume function. diff --git a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs index 38ee26c5a87..30dd915521c 100644 --- a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs +++ b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs @@ -154,7 +154,8 @@ impl<'b, 'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> Gatherer<'b, 'a, 'tcx, F> { | ty::FnDef(_, _) | ty::FnPtr(_) | ty::Dynamic(_, _, _) - | ty::Closure(_, _) + | ty::Closure(..) + | ty::CoroutineClosure(..) | ty::Coroutine(_, _) | ty::CoroutineWitness(..) | ty::Never @@ -177,7 +178,10 @@ impl<'b, 'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> Gatherer<'b, 'a, 'tcx, F> { union_path.get_or_insert(base); } } - ty::Closure(_, _) | ty::Coroutine(_, _) | ty::Tuple(_) => (), + ty::Closure(..) + | ty::CoroutineClosure(..) + | ty::Coroutine(_, _) + | ty::Tuple(_) => (), ty::Bool | ty::Char | ty::Int(_) -- cgit 1.4.1-3-g733a5 From a82bae2172499864c12a1d0b412931ad884911f7 Mon Sep 17 00:00:00 2001 From: Michael Goulet Date: Wed, 24 Jan 2024 22:27:25 +0000 Subject: Teach typeck/borrowck/solvers how to deal with async closures --- compiler/rustc_ast_lowering/src/expr.rs | 8 + .../rustc_borrowck/src/diagnostics/region_name.rs | 21 +-- compiler/rustc_borrowck/src/lib.rs | 1 + .../rustc_borrowck/src/type_check/input_output.rs | 70 +++++++- compiler/rustc_borrowck/src/type_check/mod.rs | 15 +- compiler/rustc_borrowck/src/universal_regions.rs | 58 ++++++- .../rustc_const_eval/src/transform/validate.rs | 1 + compiler/rustc_hir/src/lang_items.rs | 1 + compiler/rustc_hir_analysis/src/collect.rs | 31 ++++ compiler/rustc_hir_typeck/src/callee.rs | 85 +++++++--- compiler/rustc_hir_typeck/src/closure.rs | 87 +++++++++- compiler/rustc_hir_typeck/src/upvar.rs | 64 +++++++- compiler/rustc_middle/src/middle/lang_items.rs | 15 +- compiler/rustc_middle/src/query/mod.rs | 5 + compiler/rustc_middle/src/traits/select.rs | 7 + compiler/rustc_middle/src/ty/mod.rs | 7 +- compiler/rustc_middle/src/ty/sty.rs | 128 ++++++++++++++- compiler/rustc_mir_build/src/build/mod.rs | 1 + compiler/rustc_mir_build/src/thir/cx/expr.rs | 3 + compiler/rustc_mir_build/src/thir/cx/mod.rs | 18 +- compiler/rustc_mir_dataflow/src/value_analysis.rs | 3 + .../rustc_mir_transform/src/const_prop_lint.rs | 3 +- .../src/coverage/spans/from_mir.rs | 4 +- compiler/rustc_span/src/symbol.rs | 4 + .../src/solve/assembly/mod.rs | 18 ++ .../src/solve/assembly/structural_traits.rs | 111 ++++++++++++- .../src/solve/normalizes_to/mod.rs | 113 +++++++++++++ .../rustc_trait_selection/src/solve/trait_goals.rs | 60 +++++++ .../rustc_trait_selection/src/traits/project.rs | 182 ++++++++++++++++++++- .../src/traits/select/candidate_assembly.rs | 46 ++++++ .../src/traits/select/confirmation.rs | 50 ++++++ .../rustc_trait_selection/src/traits/select/mod.rs | 12 ++ compiler/rustc_ty_utils/src/abi.rs | 36 ++++ compiler/rustc_ty_utils/src/instance.rs | 11 ++ library/core/src/ops/async_function.rs | 8 + 35 files changed, 1221 insertions(+), 66 deletions(-) (limited to 'compiler/rustc_mir_dataflow') diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs index 2241b4cd8e5..9990d526bf7 100644 --- a/compiler/rustc_ast_lowering/src/expr.rs +++ b/compiler/rustc_ast_lowering/src/expr.rs @@ -1,3 +1,5 @@ +use std::assert_matches::assert_matches; + use super::errors::{ AsyncCoroutinesNotSupported, AwaitOnlyInAsyncFnAndBlocks, BaseExpressionDoubleDot, ClosureCannotBeStatic, CoroutineTooManyParameters, @@ -1027,6 +1029,12 @@ impl<'hir> LoweringContext<'_, 'hir> { ) -> hir::ExprKind<'hir> { let (binder_clause, generic_params) = self.lower_closure_binder(binder); + assert_matches!( + coroutine_kind, + CoroutineKind::Async { .. }, + "only async closures are supported currently" + ); + let body = self.with_new_scopes(fn_decl_span, |this| { let inner_decl = FnDecl { inputs: decl.inputs.clone(), output: FnRetTy::Default(fn_decl_span) }; diff --git a/compiler/rustc_borrowck/src/diagnostics/region_name.rs b/compiler/rustc_borrowck/src/diagnostics/region_name.rs index fc52306e10d..c91b8eaab05 100644 --- a/compiler/rustc_borrowck/src/diagnostics/region_name.rs +++ b/compiler/rustc_borrowck/src/diagnostics/region_name.rs @@ -324,9 +324,13 @@ impl<'tcx> MirBorrowckCtxt<'_, 'tcx> { ty::BoundRegionKind::BrEnv => { let def_ty = self.regioncx.universal_regions().defining_ty; - let DefiningTy::Closure(_, args) = def_ty else { - // Can't have BrEnv in functions, constants or coroutines. - bug!("BrEnv outside of closure."); + let closure_kind = match def_ty { + DefiningTy::Closure(_, args) => args.as_closure().kind(), + DefiningTy::CoroutineClosure(_, args) => args.as_coroutine_closure().kind(), + _ => { + // Can't have BrEnv in functions, constants or coroutines. + bug!("BrEnv outside of closure."); + } }; let hir::ExprKind::Closure(&hir::Closure { fn_decl_span, .. }) = tcx.hir().expect_expr(self.mir_hir_id()).kind @@ -334,21 +338,18 @@ impl<'tcx> MirBorrowckCtxt<'_, 'tcx> { bug!("Closure is not defined by a closure expr"); }; let region_name = self.synthesize_region_name(); - - let closure_kind_ty = args.as_closure().kind_ty(); - let note = match closure_kind_ty.to_opt_closure_kind() { - Some(ty::ClosureKind::Fn) => { + let note = match closure_kind { + ty::ClosureKind::Fn => { "closure implements `Fn`, so references to captured variables \ can't escape the closure" } - Some(ty::ClosureKind::FnMut) => { + ty::ClosureKind::FnMut => { "closure implements `FnMut`, so references to captured variables \ can't escape the closure" } - Some(ty::ClosureKind::FnOnce) => { + ty::ClosureKind::FnOnce => { bug!("BrEnv in a `FnOnce` closure"); } - None => bug!("Closure kind not inferred in borrow check"), }; Some(RegionName { diff --git a/compiler/rustc_borrowck/src/lib.rs b/compiler/rustc_borrowck/src/lib.rs index eaf9fc45837..bb64571889b 100644 --- a/compiler/rustc_borrowck/src/lib.rs +++ b/compiler/rustc_borrowck/src/lib.rs @@ -3,6 +3,7 @@ #![allow(internal_features)] #![feature(rustdoc_internals)] #![doc(rust_logo)] +#![feature(assert_matches)] #![feature(associated_type_bounds)] #![feature(box_patterns)] #![feature(let_chains)] diff --git a/compiler/rustc_borrowck/src/type_check/input_output.rs b/compiler/rustc_borrowck/src/type_check/input_output.rs index 59518f68ab1..a3e5088ee09 100644 --- a/compiler/rustc_borrowck/src/type_check/input_output.rs +++ b/compiler/rustc_borrowck/src/type_check/input_output.rs @@ -7,13 +7,18 @@ //! `RETURN_PLACE` the MIR arguments) are always fully normalized (and //! contain revealed `impl Trait` values). +use std::assert_matches::assert_matches; + use itertools::Itertools; -use rustc_infer::infer::BoundRegionConversionTime; +use rustc_hir as hir; +use rustc_infer::infer::type_variable::{TypeVariableOrigin, TypeVariableOriginKind}; +use rustc_infer::infer::{BoundRegionConversionTime, RegionVariableOrigin}; use rustc_middle::mir::*; use rustc_middle::ty::{self, Ty}; use rustc_span::Span; -use crate::universal_regions::UniversalRegions; +use crate::renumber::RegionCtxt; +use crate::universal_regions::{DefiningTy, UniversalRegions}; use super::{Locations, TypeChecker}; @@ -23,9 +28,11 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> { #[instrument(skip(self, body), level = "debug")] pub(super) fn check_signature_annotation(&mut self, body: &Body<'tcx>) { let mir_def_id = body.source.def_id().expect_local(); + if !self.tcx().is_closure_or_coroutine(mir_def_id.to_def_id()) { return; } + let user_provided_poly_sig = self.tcx().closure_user_provided_sig(mir_def_id); // Instantiate the canonicalized variables from user-provided signature @@ -34,12 +41,69 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> { // so that they represent the view from "inside" the closure. let user_provided_sig = self .instantiate_canonical_with_fresh_inference_vars(body.span, &user_provided_poly_sig); - let user_provided_sig = self.infcx.instantiate_binder_with_fresh_vars( + let mut user_provided_sig = self.infcx.instantiate_binder_with_fresh_vars( body.span, BoundRegionConversionTime::FnCall, user_provided_sig, ); + // FIXME(async_closures): We must apply the same transformation to our + // signature here as we do during closure checking. + if let DefiningTy::CoroutineClosure(_, args) = + self.borrowck_context.universal_regions.defining_ty + { + assert_matches!( + self.tcx().coroutine_kind(self.tcx().coroutine_for_closure(mir_def_id)), + Some(hir::CoroutineKind::Desugared( + hir::CoroutineDesugaring::Async, + hir::CoroutineSource::Closure + )), + "this needs to be modified if we're lowering non-async closures" + ); + let args = args.as_coroutine_closure(); + let tupled_upvars_ty = ty::CoroutineClosureSignature::tupled_upvars_by_closure_kind( + self.tcx(), + args.kind(), + Ty::new_tup(self.tcx(), user_provided_sig.inputs()), + args.tupled_upvars_ty(), + args.coroutine_captures_by_ref_ty(), + self.infcx.next_region_var(RegionVariableOrigin::MiscVariable(body.span), || { + RegionCtxt::Unknown + }), + ); + + let next_ty_var = || { + self.infcx.next_ty_var(TypeVariableOrigin { + span: body.span, + kind: TypeVariableOriginKind::MiscVariable, + }) + }; + let output_ty = Ty::new_coroutine( + self.tcx(), + self.tcx().coroutine_for_closure(mir_def_id), + ty::CoroutineArgs::new( + self.tcx(), + ty::CoroutineArgsParts { + parent_args: args.parent_args(), + resume_ty: next_ty_var(), + yield_ty: next_ty_var(), + witness: next_ty_var(), + return_ty: user_provided_sig.output(), + tupled_upvars_ty: tupled_upvars_ty, + }, + ) + .args, + ); + + user_provided_sig = self.tcx().mk_fn_sig( + user_provided_sig.inputs().iter().copied(), + output_ty, + user_provided_sig.c_variadic, + user_provided_sig.unsafety, + user_provided_sig.abi, + ); + } + let is_coroutine_with_implicit_resume_ty = self.tcx().is_coroutine(mir_def_id.to_def_id()) && user_provided_sig.inputs().is_empty(); diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs index 90a9844fda9..cdb187e0776 100644 --- a/compiler/rustc_borrowck/src/type_check/mod.rs +++ b/compiler/rustc_borrowck/src/type_check/mod.rs @@ -2773,15 +2773,14 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> { let typeck_root_args = ty::GenericArgs::identity_for_item(tcx, typeck_root_def_id); let parent_args = match tcx.def_kind(def_id) { + // We don't want to dispatch on 3 different kind of closures here, so take + // advantage of the fact that the `parent_args` is the same length as the + // `typeck_root_args`. DefKind::Closure => { - // FIXME(async_closures): It's kind of icky to access HIR here. - match tcx.hir_node_by_def_id(def_id).expect_closure().kind { - hir::ClosureKind::Closure => args.as_closure().parent_args(), - hir::ClosureKind::Coroutine(_) => args.as_coroutine().parent_args(), - hir::ClosureKind::CoroutineClosure(_) => { - args.as_coroutine_closure().parent_args() - } - } + // FIXME(async_closures): It may be useful to add a debug assert here + // to actually call `type_of` and check the `parent_args` are the same + // length as the `typeck_root_args`. + &args[..typeck_root_args.len()] } DefKind::InlineConst => args.as_inline_const().parent_args(), other => bug!("unexpected item {:?}", other), diff --git a/compiler/rustc_borrowck/src/universal_regions.rs b/compiler/rustc_borrowck/src/universal_regions.rs index ae8a135f090..9c266d123b3 100644 --- a/compiler/rustc_borrowck/src/universal_regions.rs +++ b/compiler/rustc_borrowck/src/universal_regions.rs @@ -97,6 +97,10 @@ pub enum DefiningTy<'tcx> { /// `ClosureArgs::coroutine_return_ty`. Coroutine(DefId, GenericArgsRef<'tcx>), + /// The MIR is a special kind of closure that returns coroutines. + /// TODO: describe how to make the sig... + CoroutineClosure(DefId, GenericArgsRef<'tcx>), + /// The MIR is a fn item with the given `DefId` and args. The signature /// of the function can be bound then with the `fn_sig` query. FnDef(DefId, GenericArgsRef<'tcx>), @@ -119,6 +123,7 @@ impl<'tcx> DefiningTy<'tcx> { pub fn upvar_tys(self) -> &'tcx ty::List> { match self { DefiningTy::Closure(_, args) => args.as_closure().upvar_tys(), + DefiningTy::CoroutineClosure(_, args) => args.as_coroutine_closure().upvar_tys(), DefiningTy::Coroutine(_, args) => args.as_coroutine().upvar_tys(), DefiningTy::FnDef(..) | DefiningTy::Const(..) | DefiningTy::InlineConst(..) => { ty::List::empty() @@ -131,7 +136,9 @@ impl<'tcx> DefiningTy<'tcx> { /// user's code. pub fn implicit_inputs(self) -> usize { match self { - DefiningTy::Closure(..) | DefiningTy::Coroutine(..) => 1, + DefiningTy::Closure(..) + | DefiningTy::CoroutineClosure(..) + | DefiningTy::Coroutine(..) => 1, DefiningTy::FnDef(..) | DefiningTy::Const(..) | DefiningTy::InlineConst(..) => 0, } } @@ -147,6 +154,7 @@ impl<'tcx> DefiningTy<'tcx> { pub fn def_id(&self) -> DefId { match *self { DefiningTy::Closure(def_id, ..) + | DefiningTy::CoroutineClosure(def_id, ..) | DefiningTy::Coroutine(def_id, ..) | DefiningTy::FnDef(def_id, ..) | DefiningTy::Const(def_id, ..) @@ -355,6 +363,9 @@ impl<'tcx> UniversalRegions<'tcx> { err.note(format!("late-bound region is {:?}", self.to_region_vid(r))); }); } + DefiningTy::CoroutineClosure(..) => { + todo!() + } DefiningTy::Coroutine(def_id, args) => { let v = with_no_trimmed_paths!( args[tcx.generics_of(def_id).parent_count..] @@ -568,6 +579,9 @@ impl<'cx, 'tcx> UniversalRegionsBuilder<'cx, 'tcx> { match *defining_ty.kind() { ty::Closure(def_id, args) => DefiningTy::Closure(def_id, args), ty::Coroutine(def_id, args) => DefiningTy::Coroutine(def_id, args), + ty::CoroutineClosure(def_id, args) => { + DefiningTy::CoroutineClosure(def_id, args) + } ty::FnDef(def_id, args) => DefiningTy::FnDef(def_id, args), _ => span_bug!( tcx.def_span(self.mir_def), @@ -623,6 +637,7 @@ impl<'cx, 'tcx> UniversalRegionsBuilder<'cx, 'tcx> { let identity_args = GenericArgs::identity_for_item(tcx, typeck_root_def_id); let fr_args = match defining_ty { DefiningTy::Closure(_, args) + | DefiningTy::CoroutineClosure(_, args) | DefiningTy::Coroutine(_, args) | DefiningTy::InlineConst(_, args) => { // In the case of closures, we rely on the fact that @@ -702,6 +717,47 @@ impl<'cx, 'tcx> UniversalRegionsBuilder<'cx, 'tcx> { ty::Binder::dummy(inputs_and_output) } + DefiningTy::CoroutineClosure(def_id, args) => { + assert_eq!(self.mir_def.to_def_id(), def_id); + let closure_sig = args.as_coroutine_closure().coroutine_closure_sig(); + let bound_vars = tcx.mk_bound_variable_kinds_from_iter( + closure_sig + .bound_vars() + .iter() + .chain(iter::once(ty::BoundVariableKind::Region(ty::BrEnv))), + ); + let br = ty::BoundRegion { + var: ty::BoundVar::from_usize(bound_vars.len() - 1), + kind: ty::BrEnv, + }; + let env_region = ty::Region::new_bound(tcx, ty::INNERMOST, br); + let closure_kind = args.as_coroutine_closure().kind(); + + let closure_ty = tcx.closure_env_ty( + Ty::new_coroutine_closure(tcx, def_id, args), + closure_kind, + env_region, + ); + + let inputs = closure_sig.skip_binder().tupled_inputs_ty.tuple_fields(); + let output = closure_sig.skip_binder().to_coroutine_given_kind_and_upvars( + tcx, + args.as_coroutine_closure().parent_args(), + tcx.coroutine_for_closure(def_id), + closure_kind, + env_region, + args.as_coroutine_closure().tupled_upvars_ty(), + args.as_coroutine_closure().coroutine_captures_by_ref_ty(), + ); + + ty::Binder::bind_with_vars( + tcx.mk_type_list_from_iter( + iter::once(closure_ty).chain(inputs).chain(iter::once(output)), + ), + bound_vars, + ) + } + DefiningTy::FnDef(def_id, _) => { let sig = tcx.fn_sig(def_id).instantiate_identity(); let sig = indices.fold_to_region_vids(tcx, sig); diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs index d3230a2455d..c4542aaa7b2 100644 --- a/compiler/rustc_const_eval/src/transform/validate.rs +++ b/compiler/rustc_const_eval/src/transform/validate.rs @@ -58,6 +58,7 @@ impl<'tcx> MirPass<'tcx> for Validator { 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: {:?} phase {:?}", body_ty, mir_phase) diff --git a/compiler/rustc_hir/src/lang_items.rs b/compiler/rustc_hir/src/lang_items.rs index 8e327c029df..9c9e72d6e65 100644 --- a/compiler/rustc_hir/src/lang_items.rs +++ b/compiler/rustc_hir/src/lang_items.rs @@ -209,6 +209,7 @@ language_item_table! { AsyncFn, sym::async_fn, async_fn_trait, Target::Trait, GenericRequirement::Exact(1); AsyncFnMut, sym::async_fn_mut, async_fn_mut_trait, Target::Trait, GenericRequirement::Exact(1); AsyncFnOnce, sym::async_fn_once, async_fn_once_trait, Target::Trait, GenericRequirement::Exact(1); + AsyncFnKindHelper, sym::async_fn_kind_helper,async_fn_kind_helper, Target::Trait, GenericRequirement::Exact(1); FnOnceOutput, sym::fn_once_output, fn_once_output, Target::AssocTy, GenericRequirement::None; diff --git a/compiler/rustc_hir_analysis/src/collect.rs b/compiler/rustc_hir_analysis/src/collect.rs index 8d862d5eb71..8fdd3cfe3cf 100644 --- a/compiler/rustc_hir_analysis/src/collect.rs +++ b/compiler/rustc_hir_analysis/src/collect.rs @@ -81,6 +81,7 @@ pub fn provide(providers: &mut Providers) { impl_trait_ref, impl_polarity, coroutine_kind, + coroutine_for_closure, collect_mod_item_types, is_type_alias_impl_trait, ..*providers @@ -1531,6 +1532,36 @@ fn coroutine_kind(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option, def_id: LocalDefId) -> DefId { + let Node::Expr(&hir::Expr { + kind: + hir::ExprKind::Closure(&rustc_hir::Closure { + kind: hir::ClosureKind::CoroutineClosure(_), + body, + .. + }), + .. + }) = tcx.hir_node_by_def_id(def_id) + else { + bug!() + }; + + let &hir::Expr { + kind: + hir::ExprKind::Closure(&rustc_hir::Closure { + def_id, + kind: hir::ClosureKind::Coroutine(_), + .. + }), + .. + } = tcx.hir().body(body).value + else { + bug!() + }; + + def_id.to_def_id() +} + fn is_type_alias_impl_trait<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> bool { match tcx.hir_node_by_def_id(def_id) { Node::Item(hir::Item { kind: hir::ItemKind::OpaqueTy(opaque), .. }) => { diff --git a/compiler/rustc_hir_typeck/src/callee.rs b/compiler/rustc_hir_typeck/src/callee.rs index 80650dcb53c..1858b2770cd 100644 --- a/compiler/rustc_hir_typeck/src/callee.rs +++ b/compiler/rustc_hir_typeck/src/callee.rs @@ -141,32 +141,67 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { return Some(CallStep::Builtin(adjusted_ty)); } - ty::Closure(def_id, args) => { + // Check whether this is a call to a closure where we + // haven't yet decided on whether the closure is fn vs + // fnmut vs fnonce. If so, we have to defer further processing. + ty::Closure(def_id, args) if self.closure_kind(adjusted_ty).is_none() => { let def_id = def_id.expect_local(); + let closure_sig = args.as_closure().sig(); + let closure_sig = self.instantiate_binder_with_fresh_vars( + call_expr.span, + infer::FnCall, + closure_sig, + ); + let adjustments = self.adjust_steps(autoderef); + self.record_deferred_call_resolution( + def_id, + DeferredCallResolution { + call_expr, + callee_expr, + closure_ty: adjusted_ty, + adjustments, + fn_sig: closure_sig, + }, + ); + return Some(CallStep::DeferredClosure(def_id, closure_sig)); + } - // Check whether this is a call to a closure where we - // haven't yet decided on whether the closure is fn vs - // fnmut vs fnonce. If so, we have to defer further processing. - if self.closure_kind(adjusted_ty).is_none() { - let closure_sig = args.as_closure().sig(); - let closure_sig = self.instantiate_binder_with_fresh_vars( - call_expr.span, - infer::FnCall, - closure_sig, - ); - let adjustments = self.adjust_steps(autoderef); - self.record_deferred_call_resolution( - def_id, - DeferredCallResolution { - call_expr, - callee_expr, - closure_ty: adjusted_ty, - adjustments, - fn_sig: closure_sig, - }, - ); - return Some(CallStep::DeferredClosure(def_id, closure_sig)); - } + ty::CoroutineClosure(def_id, args) if self.closure_kind(adjusted_ty).is_none() => { + let def_id = def_id.expect_local(); + let closure_args = args.as_coroutine_closure(); + let coroutine_closure_sig = self.instantiate_binder_with_fresh_vars( + call_expr.span, + infer::FnCall, + closure_args.coroutine_closure_sig(), + ); + let tupled_upvars_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(), + self.tcx.coroutine_for_closure(def_id), + tupled_upvars_ty, + ), + coroutine_closure_sig.c_variadic, + coroutine_closure_sig.unsafety, + coroutine_closure_sig.abi, + ); + let adjustments = self.adjust_steps(autoderef); + self.record_deferred_call_resolution( + def_id, + DeferredCallResolution { + call_expr, + callee_expr, + closure_ty: adjusted_ty, + adjustments, + fn_sig: call_sig, + }, + ); + return Some(CallStep::DeferredClosure(def_id, call_sig)); } // Hack: we know that there are traits implementing Fn for &F @@ -935,7 +970,7 @@ impl<'a, 'tcx> DeferredCallResolution<'tcx> { span_bug!( self.call_expr.span, "Expected to find a suitable `Fn`/`FnMut`/`FnOnce` implementation for `{}`", - self.adjusted_ty + self.closure_ty ) } } diff --git a/compiler/rustc_hir_typeck/src/closure.rs b/compiler/rustc_hir_typeck/src/closure.rs index 18dff5188af..1d024efdd49 100644 --- a/compiler/rustc_hir_typeck/src/closure.rs +++ b/compiler/rustc_hir_typeck/src/closure.rs @@ -63,7 +63,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { None => (None, None), }; - let ClosureSignatures { bound_sig, liberated_sig } = + let ClosureSignatures { bound_sig, mut liberated_sig } = self.sig_of_closure(expr_def_id, closure.fn_decl, closure.kind, expected_sig); debug!(?bound_sig, ?liberated_sig); @@ -125,7 +125,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Gen, _) | hir::CoroutineKind::Coroutine(_) => { let yield_ty = self.next_ty_var(TypeVariableOrigin { - kind: TypeVariableOriginKind::TypeInference, + kind: TypeVariableOriginKind::ClosureSynthetic, span: expr_span, }); self.require_type_is_sized(yield_ty, expr_span, traits::SizedYieldType); @@ -137,7 +137,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { // not a problem. hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::AsyncGen, _) => { let yield_ty = self.next_ty_var(TypeVariableOrigin { - kind: TypeVariableOriginKind::TypeInference, + kind: TypeVariableOriginKind::ClosureSynthetic, span: expr_span, }); self.require_type_is_sized(yield_ty, expr_span, traits::SizedYieldType); @@ -166,8 +166,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { let resume_ty = liberated_sig.inputs().get(0).copied().unwrap_or(tcx.types.unit); let interior = self.next_ty_var(TypeVariableOrigin { - kind: TypeVariableOriginKind::MiscVariable, - span: body.value.span, + kind: TypeVariableOriginKind::ClosureSynthetic, + span: expr_span, }); self.deferred_coroutine_interiors.borrow_mut().push(( expr_def_id, @@ -192,7 +192,82 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { Some(CoroutineTypes { resume_ty, yield_ty }), ) } - hir::ClosureKind::CoroutineClosure(_) => todo!(), + hir::ClosureKind::CoroutineClosure(kind) => { + let (bound_return_ty, bound_yield_ty) = match kind { + hir::CoroutineDesugaring::Async => { + (bound_sig.skip_binder().output(), tcx.types.unit) + } + hir::CoroutineDesugaring::Gen | hir::CoroutineDesugaring::AsyncGen => { + todo!("`gen` and `async gen` closures not supported yet") + } + }; + let resume_ty = self.next_ty_var(TypeVariableOrigin { + kind: TypeVariableOriginKind::ClosureSynthetic, + span: expr_span, + }); + let interior = self.next_ty_var(TypeVariableOrigin { + kind: TypeVariableOriginKind::ClosureSynthetic, + span: expr_span, + }); + let closure_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_captures_by_ref_ty = self.next_ty_var(TypeVariableOrigin { + kind: TypeVariableOriginKind::ClosureSynthetic, + span: expr_span, + }); + let closure_args = ty::CoroutineClosureArgs::new( + tcx, + ty::CoroutineClosureArgsParts { + parent_args, + closure_kind_ty, + signature_parts_ty: Ty::new_fn_ptr( + tcx, + bound_sig.map_bound(|sig| { + tcx.mk_fn_sig( + [ + resume_ty, + Ty::new_tup_from_iter(tcx, sig.inputs().iter().copied()), + ], + Ty::new_tup(tcx, &[bound_yield_ty, bound_return_ty]), + sig.c_variadic, + sig.unsafety, + sig.abi, + ) + }), + ), + tupled_upvars_ty, + coroutine_captures_by_ref_ty, + coroutine_witness_ty: interior, + }, + ); + + let coroutine_upvars_ty = self.next_ty_var(TypeVariableOrigin { + kind: TypeVariableOriginKind::ClosureSynthetic, + span: expr_span, + }); + liberated_sig = tcx.mk_fn_sig( + liberated_sig.inputs().iter().copied(), + tcx.liberate_late_bound_regions( + expr_def_id.to_def_id(), + closure_args.coroutine_closure_sig().map_bound(|sig| { + sig.to_coroutine( + tcx, + parent_args, + tcx.coroutine_for_closure(expr_def_id), + coroutine_upvars_ty, + ) + }), + ), + liberated_sig.c_variadic, + liberated_sig.unsafety, + liberated_sig.abi, + ); + + (Ty::new_coroutine_closure(tcx, expr_def_id.to_def_id(), closure_args.args), None) + } }; check_fn( diff --git a/compiler/rustc_hir_typeck/src/upvar.rs b/compiler/rustc_hir_typeck/src/upvar.rs index 55ffac8d92f..b087d6d9e57 100644 --- a/compiler/rustc_hir_typeck/src/upvar.rs +++ b/compiler/rustc_hir_typeck/src/upvar.rs @@ -174,6 +174,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { ty::Closure(def_id, args) => { (def_id, UpvarArgs::Closure(args), self.closure_kind(ty).is_none()) } + ty::CoroutineClosure(def_id, args) => { + (def_id, UpvarArgs::CoroutineClosure(args), self.closure_kind(ty).is_none()) + } ty::Coroutine(def_id, args) => (def_id, UpvarArgs::Coroutine(args), false), ty::Error(_) => { // #51714: skip analysis when we have already encountered type errors @@ -333,6 +336,65 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { } } + if let UpvarArgs::CoroutineClosure(args) = args { + let closure_env_region: ty::Region<'_> = ty::Region::new_bound( + self.tcx, + ty::INNERMOST, + ty::BoundRegion { + var: ty::BoundVar::from_usize(0), + kind: ty::BoundRegionKind::BrEnv, + }, + ); + let tupled_upvars_ty_for_borrow = Ty::new_tup_from_iter( + self.tcx, + self.typeck_results + .borrow() + .closure_min_captures_flattened( + self.tcx.coroutine_for_closure(closure_def_id).expect_local(), + ) + // Skip the captures that are just moving the closure's args + // into the coroutine. These are always by move. + .skip( + args.as_coroutine_closure() + .coroutine_closure_sig() + .skip_binder() + .tupled_inputs_ty + .tuple_fields() + .len(), + ) + .map(|captured_place| { + let upvar_ty = captured_place.place.ty(); + let capture = captured_place.info.capture_kind; + apply_capture_kind_on_capture_ty( + self.tcx, + upvar_ty, + capture, + Some(closure_env_region), + ) + }), + ); + let coroutine_captures_by_ref_ty = Ty::new_fn_ptr( + self.tcx, + ty::Binder::bind_with_vars( + self.tcx.mk_fn_sig( + [], + tupled_upvars_ty_for_borrow, + false, + hir::Unsafety::Normal, + rustc_target::spec::abi::Abi::Rust, + ), + self.tcx.mk_bound_variable_kinds(&[ty::BoundVariableKind::Region( + ty::BoundRegionKind::BrEnv, + )]), + ), + ); + self.demand_eqtype( + span, + args.as_coroutine_closure().coroutine_captures_by_ref_ty(), + coroutine_captures_by_ref_ty, + ); + } + self.log_closure_min_capture_info(closure_def_id, span); // Now that we've analyzed the closure, we know how each @@ -602,7 +664,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { }; // Go through each entry in the current list of min_captures - // - if ancestor is found, update it's capture kind to account for current place's + // - if ancestor is found, update its capture kind to account for current place's // capture information. // // - if descendant is found, remove it from the list, and update the current place's diff --git a/compiler/rustc_middle/src/middle/lang_items.rs b/compiler/rustc_middle/src/middle/lang_items.rs index f92c72c8a58..a4e193ba2c9 100644 --- a/compiler/rustc_middle/src/middle/lang_items.rs +++ b/compiler/rustc_middle/src/middle/lang_items.rs @@ -23,7 +23,7 @@ impl<'tcx> TyCtxt<'tcx> { }) } - /// Given a [`DefId`] of a [`Fn`], [`FnMut`] or [`FnOnce`] traits, + /// Given a [`DefId`] of one of the [`Fn`], [`FnMut`] or [`FnOnce`] traits, /// returns a corresponding [`ty::ClosureKind`]. /// For any other [`DefId`] return `None`. pub fn fn_trait_kind_from_def_id(self, id: DefId) -> Option { @@ -36,6 +36,19 @@ impl<'tcx> TyCtxt<'tcx> { } } + /// Given a [`DefId`] of one of the `AsyncFn`, `AsyncFnMut` or `AsyncFnOnce` traits, + /// returns a corresponding [`ty::ClosureKind`]. + /// For any other [`DefId`] return `None`. + pub fn async_fn_trait_kind_from_def_id(self, id: DefId) -> Option { + let items = self.lang_items(); + match Some(id) { + x if x == items.async_fn_trait() => Some(ty::ClosureKind::Fn), + x if x == items.async_fn_mut_trait() => Some(ty::ClosureKind::FnMut), + x if x == items.async_fn_once_trait() => Some(ty::ClosureKind::FnOnce), + _ => None, + } + } + /// Given a [`ty::ClosureKind`], get the [`DefId`] of its corresponding `Fn`-family /// trait, if it is defined. pub fn fn_trait_kind_to_def_id(self, kind: ty::ClosureKind) -> Option { diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs index 2438f826441..f9ab32b16f5 100644 --- a/compiler/rustc_middle/src/query/mod.rs +++ b/compiler/rustc_middle/src/query/mod.rs @@ -755,6 +755,11 @@ rustc_queries! { separate_provide_extern } + query coroutine_for_closure(def_id: DefId) -> DefId { + desc { |_tcx| "TODO" } + separate_provide_extern + } + /// Gets a map with the variance of every item; use `variances_of` instead. query crate_variances(_: ()) -> &'tcx ty::CrateVariancesMap<'tcx> { arena_cache diff --git a/compiler/rustc_middle/src/traits/select.rs b/compiler/rustc_middle/src/traits/select.rs index 64f4af08e12..4e11575cf98 100644 --- a/compiler/rustc_middle/src/traits/select.rs +++ b/compiler/rustc_middle/src/traits/select.rs @@ -134,6 +134,13 @@ pub enum SelectionCandidate<'tcx> { is_const: bool, }, + /// Implementation of an `AsyncFn`-family trait by one of the anonymous types + /// generated for an `async ||` expression. + AsyncClosureCandidate, + + // TODO: + AsyncFnKindHelperCandidate, + /// Implementation of a `Coroutine` trait by one of the anonymous types /// generated for a coroutine. CoroutineCandidate, diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs index 8e51db82f95..4348d01eff5 100644 --- a/compiler/rustc_middle/src/ty/mod.rs +++ b/compiler/rustc_middle/src/ty/mod.rs @@ -105,9 +105,10 @@ pub use self::region::{ pub use self::rvalue_scopes::RvalueScopes; pub use self::sty::{ AliasTy, Article, Binder, BoundTy, BoundTyKind, BoundVariableKind, CanonicalPolyFnSig, - ClosureArgs, ClosureArgsParts, CoroutineArgs, CoroutineArgsParts, FnSig, GenSig, - InlineConstArgs, InlineConstArgsParts, ParamConst, ParamTy, PolyFnSig, TyKind, TypeAndMut, - UpvarArgs, VarianceDiagInfo, + ClosureArgs, ClosureArgsParts, CoroutineArgs, CoroutineArgsParts, CoroutineClosureArgs, + CoroutineClosureArgsParts, CoroutineClosureSignature, FnSig, GenSig, InlineConstArgs, + InlineConstArgsParts, ParamConst, ParamTy, PolyFnSig, TyKind, TypeAndMut, UpvarArgs, + VarianceDiagInfo, }; pub use self::trait_def::TraitDef; pub use self::typeck_results::{ diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs index c047f6a0521..e2a2e24f06d 100644 --- a/compiler/rustc_middle/src/ty/sty.rs +++ b/compiler/rustc_middle/src/ty/sty.rs @@ -36,6 +36,7 @@ use rustc_type_ir::TyKind as IrTyKind; use rustc_type_ir::TyKind::*; use rustc_type_ir::TypeAndMut as IrTypeAndMut; +use super::fold::FnMutDelegate; use super::GenericParamDefKind; // Re-export and re-parameterize some `I = TyCtxt<'tcx>` types here @@ -351,6 +352,27 @@ impl<'tcx> CoroutineClosureArgs<'tcx> { self.split().signature_parts_ty } + pub fn coroutine_closure_sig(self) -> ty::Binder<'tcx, CoroutineClosureSignature<'tcx>> { + let interior = self.coroutine_witness_ty(); + let ty::FnPtr(sig) = self.signature_parts_ty().kind() else { bug!() }; + sig.map_bound(|sig| { + let [resume_ty, tupled_inputs_ty] = *sig.inputs() else { + bug!(); + }; + let [yield_ty, return_ty] = **sig.output().tuple_fields() else { bug!() }; + CoroutineClosureSignature { + interior, + tupled_inputs_ty, + resume_ty, + yield_ty, + return_ty, + c_variadic: sig.c_variadic, + unsafety: sig.unsafety, + abi: sig.abi, + } + }) + } + pub fn coroutine_captures_by_ref_ty(self) -> Ty<'tcx> { self.split().coroutine_captures_by_ref_ty } @@ -360,6 +382,103 @@ impl<'tcx> CoroutineClosureArgs<'tcx> { } } +#[derive(Copy, Clone, PartialEq, Eq, Debug, TypeFoldable, TypeVisitable)] +pub struct CoroutineClosureSignature<'tcx> { + pub interior: Ty<'tcx>, + pub tupled_inputs_ty: Ty<'tcx>, + pub resume_ty: Ty<'tcx>, + pub yield_ty: Ty<'tcx>, + pub return_ty: Ty<'tcx>, + pub c_variadic: bool, + pub unsafety: hir::Unsafety, + pub abi: abi::Abi, +} + +impl<'tcx> CoroutineClosureSignature<'tcx> { + pub fn to_coroutine( + self, + tcx: TyCtxt<'tcx>, + parent_args: &'tcx [GenericArg<'tcx>], + coroutine_def_id: DefId, + tupled_upvars_ty: Ty<'tcx>, + ) -> Ty<'tcx> { + let coroutine_args = ty::CoroutineArgs::new( + tcx, + ty::CoroutineArgsParts { + parent_args, + resume_ty: self.resume_ty, + yield_ty: self.yield_ty, + return_ty: self.return_ty, + witness: self.interior, + tupled_upvars_ty, + }, + ); + + Ty::new_coroutine(tcx, coroutine_def_id, coroutine_args.args) + } + + pub fn to_coroutine_given_kind_and_upvars( + self, + tcx: TyCtxt<'tcx>, + parent_args: &'tcx [GenericArg<'tcx>], + coroutine_def_id: DefId, + closure_kind: ty::ClosureKind, + env_region: ty::Region<'tcx>, + closure_tupled_upvars_ty: Ty<'tcx>, + coroutine_captures_by_ref_ty: Ty<'tcx>, + ) -> Ty<'tcx> { + let tupled_upvars_ty = Self::tupled_upvars_by_closure_kind( + tcx, + closure_kind, + self.tupled_inputs_ty, + closure_tupled_upvars_ty, + coroutine_captures_by_ref_ty, + env_region, + ); + + self.to_coroutine(tcx, parent_args, coroutine_def_id, tupled_upvars_ty) + } + + /// Given a closure kind, compute the tupled upvars that the given coroutine would return. + pub fn tupled_upvars_by_closure_kind( + tcx: TyCtxt<'tcx>, + kind: ty::ClosureKind, + tupled_inputs_ty: Ty<'tcx>, + closure_tupled_upvars_ty: Ty<'tcx>, + coroutine_captures_by_ref_ty: Ty<'tcx>, + env_region: ty::Region<'tcx>, + ) -> Ty<'tcx> { + match kind { + ty::ClosureKind::Fn | ty::ClosureKind::FnMut => { + let ty::FnPtr(sig) = *coroutine_captures_by_ref_ty.kind() else { + bug!(); + }; + let coroutine_captures_by_ref_ty = tcx.replace_escaping_bound_vars_uncached( + sig.output().skip_binder(), + FnMutDelegate { + consts: &mut |c, t| ty::Const::new_bound(tcx, ty::INNERMOST, c, t), + types: &mut |t| Ty::new_bound(tcx, ty::INNERMOST, t), + regions: &mut |_| env_region, + }, + ); + Ty::new_tup_from_iter( + tcx, + tupled_inputs_ty + .tuple_fields() + .iter() + .chain(coroutine_captures_by_ref_ty.tuple_fields()), + ) + } + ty::ClosureKind::FnOnce => Ty::new_tup_from_iter( + tcx, + tupled_inputs_ty + .tuple_fields() + .iter() + .chain(closure_tupled_upvars_ty.tuple_fields()), + ), + } + } +} /// Similar to `ClosureArgs`; see the above documentation for more. #[derive(Copy, Clone, PartialEq, Eq, Debug, TypeFoldable, TypeVisitable)] pub struct CoroutineArgs<'tcx> { @@ -1495,7 +1614,7 @@ impl<'tcx> Ty<'tcx> { ) -> Ty<'tcx> { debug_assert_eq!( closure_args.len(), - tcx.generics_of(tcx.typeck_root_def_id(def_id)).count() + 3, + tcx.generics_of(tcx.typeck_root_def_id(def_id)).count() + 5, "closure constructed with incorrect substitutions" ); Ty::new(tcx, CoroutineClosure(def_id, closure_args)) @@ -1835,6 +1954,11 @@ impl<'tcx> Ty<'tcx> { matches!(self.kind(), Coroutine(..)) } + #[inline] + pub fn is_coroutine_closure(self) -> bool { + matches!(self.kind(), CoroutineClosure(..)) + } + #[inline] pub fn is_integral(self) -> bool { matches!(self.kind(), Infer(IntVar(_)) | Int(_) | Uint(_)) @@ -2144,7 +2268,7 @@ impl<'tcx> Ty<'tcx> { // "Bound" types appear in canonical queries when the // closure type is not yet known - Bound(..) | Infer(_) => None, + Bound(..) | Param(_) | Infer(_) => None, Error(_) => Some(ty::ClosureKind::Fn), diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs index 9bcfe9fbc33..9a2f2ceced1 100644 --- a/compiler/rustc_mir_build/src/build/mod.rs +++ b/compiler/rustc_mir_build/src/build/mod.rs @@ -822,6 +822,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> { let upvar_args = match closure_ty.kind() { ty::Closure(_, args) => ty::UpvarArgs::Closure(args), ty::Coroutine(_, args) => ty::UpvarArgs::Coroutine(args), + ty::CoroutineClosure(_, args) => ty::UpvarArgs::CoroutineClosure(args), _ => return, }; diff --git a/compiler/rustc_mir_build/src/thir/cx/expr.rs b/compiler/rustc_mir_build/src/thir/cx/expr.rs index 22094c112fc..35adcc33480 100644 --- a/compiler/rustc_mir_build/src/thir/cx/expr.rs +++ b/compiler/rustc_mir_build/src/thir/cx/expr.rs @@ -556,6 +556,9 @@ impl<'tcx> Cx<'tcx> { ty::Coroutine(def_id, args) => { (def_id, UpvarArgs::Coroutine(args), Some(tcx.coroutine_movability(def_id))) } + ty::CoroutineClosure(def_id, args) => { + (def_id, UpvarArgs::CoroutineClosure(args), None) + } _ => { span_bug!(expr.span, "closure expr w/o closure type: {:?}", closure_ty); } diff --git a/compiler/rustc_mir_build/src/thir/cx/mod.rs b/compiler/rustc_mir_build/src/thir/cx/mod.rs index 5d0bb3954cc..f4f591d15b9 100644 --- a/compiler/rustc_mir_build/src/thir/cx/mod.rs +++ b/compiler/rustc_mir_build/src/thir/cx/mod.rs @@ -127,10 +127,24 @@ impl<'tcx> Cx<'tcx> { ty::Coroutine(..) => { Param { ty: closure_ty, pat: None, ty_span: None, self_kind: None, hir_id: None } } - ty::Closure(_, closure_args) => { + ty::Closure(_, args) => { let closure_env_ty = self.tcx.closure_env_ty( closure_ty, - closure_args.as_closure().kind(), + args.as_closure().kind(), + self.tcx.lifetimes.re_erased, + ); + Param { + ty: closure_env_ty, + pat: None, + ty_span: None, + self_kind: None, + hir_id: None, + } + } + ty::CoroutineClosure(_, args) => { + let closure_env_ty = self.tcx.closure_env_ty( + closure_ty, + args.as_coroutine_closure().kind(), self.tcx.lifetimes.re_erased, ); Param { diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs index 2802f5491be..a85e53ff241 100644 --- a/compiler/rustc_mir_dataflow/src/value_analysis.rs +++ b/compiler/rustc_mir_dataflow/src/value_analysis.rs @@ -1149,6 +1149,9 @@ pub fn iter_fields<'tcx>( ty::Closure(_, args) => { iter_fields(args.as_closure().tupled_upvars_ty(), tcx, param_env, f); } + ty::CoroutineClosure(_, args) => { + iter_fields(args.as_coroutine_closure().tupled_upvars_ty(), tcx, param_env, f); + } _ => (), } } 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/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_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index dbfc89c2d49..587eb1c7cc5 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -167,6 +167,9 @@ symbols! { Break, C, CStr, + CallFuture, + CallMutFuture, + CallOnceFuture, Capture, Center, Cleanup, @@ -420,6 +423,7 @@ symbols! { async_closure, async_fn, async_fn_in_trait, + async_fn_kind_helper, async_fn_mut, async_fn_once, async_fn_track_caller, diff --git a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs index e1c68039e79..8451fbcc434 100644 --- a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs +++ b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs @@ -182,6 +182,20 @@ pub(super) trait GoalKind<'tcx>: kind: ty::ClosureKind, ) -> QueryResult<'tcx>; + /// An async closure is known to implement the `AsyncFn` family of traits + /// where `A` is given by the signature of the type. + fn consider_builtin_async_fn_trait_candidates( + ecx: &mut EvalCtxt<'_, 'tcx>, + goal: Goal<'tcx, Self>, + kind: ty::ClosureKind, + ) -> QueryResult<'tcx>; + + /// TODO: + fn consider_builtin_async_fn_kind_helper_candidate( + ecx: &mut EvalCtxt<'_, 'tcx>, + goal: Goal<'tcx, Self>, + ) -> QueryResult<'tcx>; + /// `Tuple` is implemented if the `Self` type is a tuple. fn consider_builtin_tuple_candidate( ecx: &mut EvalCtxt<'_, 'tcx>, @@ -461,6 +475,10 @@ impl<'tcx> EvalCtxt<'_, 'tcx> { G::consider_builtin_fn_ptr_trait_candidate(self, goal) } else if let Some(kind) = self.tcx().fn_trait_kind_from_def_id(trait_def_id) { G::consider_builtin_fn_trait_candidates(self, goal, kind) + } else if let Some(kind) = self.tcx().async_fn_trait_kind_from_def_id(trait_def_id) { + G::consider_builtin_async_fn_trait_candidates(self, goal, kind) + } else if lang_items.async_fn_kind_helper() == Some(trait_def_id) { + G::consider_builtin_async_fn_kind_helper_candidate(self, goal) } else if lang_items.tuple_trait() == Some(trait_def_id) { G::consider_builtin_tuple_candidate(self, goal) } else if lang_items.pointee_trait() == Some(trait_def_id) { diff --git a/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs b/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs index 3c571e1d96f..c35134c78eb 100644 --- a/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs +++ b/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs @@ -1,11 +1,12 @@ //! Code which is used by built-in goals that match "structurally", such a auto //! traits, `Copy`/`Clone`. use rustc_data_structures::fx::FxHashMap; +use rustc_hir::LangItem; use rustc_hir::{def_id::DefId, Movability, Mutability}; use rustc_infer::traits::query::NoSolution; use rustc_middle::traits::solve::Goal; use rustc_middle::ty::{ - self, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt, + self, ToPredicate, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt, }; use crate::solve::EvalCtxt; @@ -306,6 +307,114 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_callable<'tcx>( } } +// Returns a binder of the tupled inputs types, output type, and coroutine type +// from a builtin async closure type. +pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tcx>( + tcx: TyCtxt<'tcx>, + self_ty: Ty<'tcx>, + goal_kind: ty::ClosureKind, + env_region: ty::Region<'tcx>, +) -> Result<(ty::Binder<'tcx, (Ty<'tcx>, Ty<'tcx>, Ty<'tcx>)>, Vec>), NoSolution> +{ + match *self_ty.kind() { + ty::CoroutineClosure(def_id, args) => { + let args = args.as_coroutine_closure(); + let kind_ty = args.kind_ty(); + + if let Some(closure_kind) = kind_ty.to_opt_closure_kind() { + if !closure_kind.extends(goal_kind) { + return Err(NoSolution); + } + Ok(( + args.coroutine_closure_sig().map_bound(|sig| { + let coroutine_ty = sig.to_coroutine_given_kind_and_upvars( + tcx, + args.parent_args(), + tcx.coroutine_for_closure(def_id), + goal_kind, + env_region, + args.tupled_upvars_ty(), + args.coroutine_captures_by_ref_ty(), + ); + (sig.tupled_inputs_ty, sig.return_ty, coroutine_ty) + }), + vec![], + )) + } else { + let helper_trait_def_id = tcx.require_lang_item(LangItem::AsyncFnKindHelper, None); + // FIXME(async_closures): Make this into a lang item. + let upvars_projection_def_id = tcx + .associated_items(helper_trait_def_id) + .in_definition_order() + .next() + .unwrap() + .def_id; + Ok(( + args.coroutine_closure_sig().map_bound(|sig| { + let tupled_upvars_ty = Ty::new_projection( + tcx, + upvars_projection_def_id, + [ + ty::GenericArg::from(kind_ty), + Ty::from_closure_kind(tcx, goal_kind).into(), + env_region.into(), + sig.tupled_inputs_ty.into(), + args.tupled_upvars_ty().into(), + args.coroutine_captures_by_ref_ty().into(), + ], + ); + let coroutine_ty = sig.to_coroutine( + tcx, + args.parent_args(), + tcx.coroutine_for_closure(def_id), + tupled_upvars_ty, + ); + (sig.tupled_inputs_ty, sig.return_ty, coroutine_ty) + }), + vec![ + ty::TraitRef::new( + tcx, + helper_trait_def_id, + [kind_ty, Ty::from_closure_kind(tcx, goal_kind)], + ) + .to_predicate(tcx), + ], + )) + } + } + + ty::FnDef(..) | ty::FnPtr(..) | ty::Closure(..) => Err(NoSolution), + + ty::Bool + | ty::Char + | ty::Int(_) + | ty::Uint(_) + | ty::Float(_) + | ty::Adt(_, _) + | ty::Foreign(_) + | ty::Str + | ty::Array(_, _) + | ty::Slice(_) + | ty::RawPtr(_) + | ty::Ref(_, _, _) + | ty::Dynamic(_, _, _) + | ty::Coroutine(_, _) + | ty::CoroutineWitness(..) + | ty::Never + | ty::Tuple(_) + | ty::Alias(_, _) + | ty::Param(_) + | ty::Placeholder(..) + | ty::Infer(ty::IntVar(_) | ty::FloatVar(_)) + | ty::Error(_) => Err(NoSolution), + + ty::Bound(..) + | ty::Infer(ty::TyVar(_) | ty::FreshTy(_) | ty::FreshIntTy(_) | ty::FreshFloatTy(_)) => { + bug!("unexpected type `{self_ty}`") + } + } +} + /// Assemble a list of predicates that would be present on a theoretical /// user impl for an object type. These predicates must be checked any time /// we assemble a built-in object candidate for an object type, since they diff --git a/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs b/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs index e1c6f67f05e..47ba549022d 100644 --- a/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs +++ b/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs @@ -366,6 +366,119 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> { Self::consider_implied_clause(ecx, goal, pred, [goal.with(tcx, output_is_sized_pred)]) } + fn consider_builtin_async_fn_trait_candidates( + ecx: &mut EvalCtxt<'_, 'tcx>, + goal: Goal<'tcx, Self>, + goal_kind: ty::ClosureKind, + ) -> QueryResult<'tcx> { + let tcx = ecx.tcx(); + + let env_region = match goal_kind { + ty::ClosureKind::Fn | ty::ClosureKind::FnMut => goal.predicate.alias.args.region_at(2), + // Doesn't matter what this region is + ty::ClosureKind::FnOnce => tcx.lifetimes.re_static, + }; + let (tupled_inputs_and_output_and_coroutine, nested_preds) = + structural_traits::extract_tupled_inputs_and_output_from_async_callable( + tcx, + goal.predicate.self_ty(), + goal_kind, + env_region, + )?; + let output_is_sized_pred = + tupled_inputs_and_output_and_coroutine.map_bound(|(_, output, _)| { + ty::TraitRef::from_lang_item(tcx, LangItem::Sized, DUMMY_SP, [output]) + }); + + let pred = tupled_inputs_and_output_and_coroutine + .map_bound(|(inputs, output, coroutine)| { + let (projection_ty, term) = match tcx.item_name(goal.predicate.def_id()) { + sym::CallOnceFuture => ( + ty::AliasTy::new( + tcx, + goal.predicate.def_id(), + [goal.predicate.self_ty(), inputs], + ), + coroutine.into(), + ), + sym::CallMutFuture | sym::CallFuture => ( + ty::AliasTy::new( + tcx, + goal.predicate.def_id(), + [ + ty::GenericArg::from(goal.predicate.self_ty()), + inputs.into(), + env_region.into(), + ], + ), + coroutine.into(), + ), + sym::Output => ( + ty::AliasTy::new( + tcx, + goal.predicate.def_id(), + [ty::GenericArg::from(goal.predicate.self_ty()), inputs.into()], + ), + output.into(), + ), + name => bug!("no such associated type: {name}"), + }; + ty::ProjectionPredicate { projection_ty, term } + }) + .to_predicate(tcx); + + // A built-in `AsyncFn` impl only holds if the output is sized. + // (FIXME: technically we only need to check this if the type is a fn ptr...) + Self::consider_implied_clause( + ecx, + goal, + pred, + [goal.with(tcx, output_is_sized_pred)] + .into_iter() + .chain(nested_preds.into_iter().map(|pred| goal.with(tcx, pred))), + ) + } + + fn consider_builtin_async_fn_kind_helper_candidate( + ecx: &mut EvalCtxt<'_, 'tcx>, + goal: Goal<'tcx, Self>, + ) -> QueryResult<'tcx> { + let [ + closure_fn_kind_ty, + goal_kind_ty, + borrow_region, + tupled_inputs_ty, + tupled_upvars_ty, + coroutine_captures_by_ref_ty, + ] = **goal.predicate.alias.args + else { + bug!(); + }; + + let Some(closure_kind) = closure_fn_kind_ty.expect_ty().to_opt_closure_kind() else { + // We don't need to worry about the self type being an infer var. + return Err(NoSolution); + }; + let Some(goal_kind) = goal_kind_ty.expect_ty().to_opt_closure_kind() else { + return Err(NoSolution); + }; + if !closure_kind.extends(goal_kind) { + return Err(NoSolution); + } + + let upvars_ty = ty::CoroutineClosureSignature::tupled_upvars_by_closure_kind( + ecx.tcx(), + goal_kind, + tupled_inputs_ty.expect_ty(), + tupled_upvars_ty.expect_ty(), + coroutine_captures_by_ref_ty.expect_ty(), + borrow_region.expect_region(), + ); + + ecx.eq(goal.param_env, goal.predicate.term.ty().unwrap(), upvars_ty)?; + ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes) + } + fn consider_builtin_tuple_candidate( _ecx: &mut EvalCtxt<'_, 'tcx>, goal: Goal<'tcx, Self>, diff --git a/compiler/rustc_trait_selection/src/solve/trait_goals.rs b/compiler/rustc_trait_selection/src/solve/trait_goals.rs index efaad47b6dd..fd09a6b671d 100644 --- a/compiler/rustc_trait_selection/src/solve/trait_goals.rs +++ b/compiler/rustc_trait_selection/src/solve/trait_goals.rs @@ -303,6 +303,66 @@ impl<'tcx> assembly::GoalKind<'tcx> for TraitPredicate<'tcx> { Self::consider_implied_clause(ecx, goal, pred, [goal.with(tcx, output_is_sized_pred)]) } + fn consider_builtin_async_fn_trait_candidates( + ecx: &mut EvalCtxt<'_, 'tcx>, + goal: Goal<'tcx, Self>, + goal_kind: ty::ClosureKind, + ) -> QueryResult<'tcx> { + if goal.predicate.polarity != ty::ImplPolarity::Positive { + return Err(NoSolution); + } + + let tcx = ecx.tcx(); + let (tupled_inputs_and_output_and_coroutine, nested_preds) = + structural_traits::extract_tupled_inputs_and_output_from_async_callable( + tcx, + goal.predicate.self_ty(), + goal_kind, + // This region doesn't matter because we're throwing away the coroutine type + tcx.lifetimes.re_static, + )?; + let output_is_sized_pred = + tupled_inputs_and_output_and_coroutine.map_bound(|(_, output, _)| { + ty::TraitRef::from_lang_item(tcx, LangItem::Sized, DUMMY_SP, [output]) + }); + + let pred = tupled_inputs_and_output_and_coroutine + .map_bound(|(inputs, _, _)| { + ty::TraitRef::new(tcx, goal.predicate.def_id(), [goal.predicate.self_ty(), inputs]) + }) + .to_predicate(tcx); + // A built-in `AsyncFn` impl only holds if the output is sized. + // (FIXME: technically we only need to check this if the type is a fn ptr...) + Self::consider_implied_clause( + ecx, + goal, + pred, + [goal.with(tcx, output_is_sized_pred)] + .into_iter() + .chain(nested_preds.into_iter().map(|pred| goal.with(tcx, pred))), + ) + } + + fn consider_builtin_async_fn_kind_helper_candidate( + ecx: &mut EvalCtxt<'_, 'tcx>, + goal: Goal<'tcx, Self>, + ) -> QueryResult<'tcx> { + let [closure_fn_kind_ty, goal_kind_ty] = **goal.predicate.trait_ref.args else { + bug!(); + }; + + let Some(closure_kind) = closure_fn_kind_ty.expect_ty().to_opt_closure_kind() else { + // We don't need to worry about the self type being an infer var. + return Err(NoSolution); + }; + let goal_kind = goal_kind_ty.expect_ty().to_opt_closure_kind().unwrap(); + if closure_kind.extends(goal_kind) { + ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes) + } else { + Err(NoSolution) + } + } + fn consider_builtin_tuple_candidate( ecx: &mut EvalCtxt<'_, 'tcx>, goal: Goal<'tcx, Self>, diff --git a/compiler/rustc_trait_selection/src/traits/project.rs b/compiler/rustc_trait_selection/src/traits/project.rs index a960befcd4b..648f14beaa7 100644 --- a/compiler/rustc_trait_selection/src/traits/project.rs +++ b/compiler/rustc_trait_selection/src/traits/project.rs @@ -1833,10 +1833,28 @@ fn assemble_candidates_from_impls<'cx, 'tcx>( lang_items.fn_trait(), lang_items.fn_mut_trait(), lang_items.fn_once_trait(), + lang_items.async_fn_trait(), + lang_items.async_fn_mut_trait(), + lang_items.async_fn_once_trait(), ].contains(&Some(trait_ref.def_id)) { true - }else if lang_items.discriminant_kind_trait() == Some(trait_ref.def_id) { + } else if lang_items.async_fn_kind_helper() == Some(trait_ref.def_id) { + // FIXME(async_closures): Validity constraints here could be cleaned up. + if obligation.predicate.args.type_at(0).is_ty_var() + || obligation.predicate.args.type_at(4).is_ty_var() + || obligation.predicate.args.type_at(5).is_ty_var() + { + candidate_set.mark_ambiguous(); + true + } else if obligation.predicate.args.type_at(0).to_opt_closure_kind().is_some() + && obligation.predicate.args.type_at(1).to_opt_closure_kind().is_some() + { + true + } else { + false + } + } else if lang_items.discriminant_kind_trait() == Some(trait_ref.def_id) { match self_ty.kind() { ty::Bool | ty::Char @@ -2061,6 +2079,10 @@ fn confirm_select_candidate<'cx, 'tcx>( } else { confirm_fn_pointer_candidate(selcx, obligation, data) } + } else if selcx.tcx().async_fn_trait_kind_from_def_id(trait_def_id).is_some() { + confirm_async_closure_candidate(selcx, obligation, data) + } else if lang_items.async_fn_kind_helper() == Some(trait_def_id) { + confirm_async_fn_kind_helper_candidate(selcx, obligation, data) } else { confirm_builtin_candidate(selcx, obligation, data) } @@ -2421,6 +2443,164 @@ fn confirm_callable_candidate<'cx, 'tcx>( confirm_param_env_candidate(selcx, obligation, predicate, true) } +fn confirm_async_closure_candidate<'cx, 'tcx>( + selcx: &mut SelectionContext<'cx, 'tcx>, + obligation: &ProjectionTyObligation<'tcx>, + mut nested: Vec>, +) -> Progress<'tcx> { + let self_ty = selcx.infcx.shallow_resolve(obligation.predicate.self_ty()); + let ty::CoroutineClosure(def_id, args) = *self_ty.kind() else { + unreachable!( + "expected coroutine-closure self type for coroutine-closure candidate, found {self_ty}" + ) + }; + let args = args.as_coroutine_closure(); + let kind_ty = args.kind_ty(); + + let tcx = selcx.tcx(); + let goal_kind = + tcx.async_fn_trait_kind_from_def_id(obligation.predicate.trait_def_id(tcx)).unwrap(); + + let helper_trait_def_id = tcx.require_lang_item(LangItem::AsyncFnKindHelper, None); + nested.push(obligation.with( + tcx, + ty::TraitRef::new( + tcx, + helper_trait_def_id, + [kind_ty, Ty::from_closure_kind(tcx, goal_kind)], + ), + )); + + let env_region = match goal_kind { + ty::ClosureKind::Fn | ty::ClosureKind::FnMut => obligation.predicate.args.region_at(2), + ty::ClosureKind::FnOnce => tcx.lifetimes.re_static, + }; + + // FIXME(async_closures): Make this into a lang item. + let upvars_projection_def_id = + tcx.associated_items(helper_trait_def_id).in_definition_order().next().unwrap().def_id; + + // FIXME(async_closures): Confirmation is kind of a mess here. Ideally, + // we'd short-circuit when we know that the goal_kind >= closure_kind, and not + // register a nested predicate or create a new projection ty here. But I'm too + // lazy to make this more efficient atm, and we can always tweak it later, + // since all this does is make the solver do more work. + // + // The code duplication due to the different length args is kind of weird, too. + let poly_cache_entry = args.coroutine_closure_sig().map_bound(|sig| { + let (projection_ty, term) = match tcx.item_name(obligation.predicate.def_id) { + sym::CallOnceFuture => { + let tupled_upvars_ty = Ty::new_projection( + tcx, + upvars_projection_def_id, + [ + ty::GenericArg::from(kind_ty), + Ty::from_closure_kind(tcx, goal_kind).into(), + env_region.into(), + sig.tupled_inputs_ty.into(), + args.tupled_upvars_ty().into(), + args.coroutine_captures_by_ref_ty().into(), + ], + ); + let coroutine_ty = sig.to_coroutine( + tcx, + args.parent_args(), + tcx.coroutine_for_closure(def_id), + tupled_upvars_ty, + ); + ( + ty::AliasTy::new( + tcx, + obligation.predicate.def_id, + [self_ty, sig.tupled_inputs_ty], + ), + coroutine_ty.into(), + ) + } + sym::CallMutFuture | sym::CallFuture => { + let tupled_upvars_ty = Ty::new_projection( + tcx, + upvars_projection_def_id, + [ + ty::GenericArg::from(kind_ty), + Ty::from_closure_kind(tcx, goal_kind).into(), + env_region.into(), + sig.tupled_inputs_ty.into(), + args.tupled_upvars_ty().into(), + args.coroutine_captures_by_ref_ty().into(), + ], + ); + let coroutine_ty = sig.to_coroutine( + tcx, + args.parent_args(), + tcx.coroutine_for_closure(def_id), + tupled_upvars_ty, + ); + ( + ty::AliasTy::new( + tcx, + obligation.predicate.def_id, + [ + ty::GenericArg::from(self_ty), + sig.tupled_inputs_ty.into(), + env_region.into(), + ], + ), + coroutine_ty.into(), + ) + } + sym::Output => ( + ty::AliasTy::new(tcx, obligation.predicate.def_id, [self_ty, sig.tupled_inputs_ty]), + sig.return_ty.into(), + ), + name => bug!("no such associated type: {name}"), + }; + ty::ProjectionPredicate { projection_ty, term } + }); + + confirm_param_env_candidate(selcx, obligation, poly_cache_entry, true) + .with_addl_obligations(nested) +} + +fn confirm_async_fn_kind_helper_candidate<'cx, 'tcx>( + selcx: &mut SelectionContext<'cx, 'tcx>, + obligation: &ProjectionTyObligation<'tcx>, + nested: Vec>, +) -> Progress<'tcx> { + let [ + // We already checked that the goal_kind >= closure_kind + _closure_kind_ty, + goal_kind_ty, + borrow_region, + tupled_inputs_ty, + tupled_upvars_ty, + coroutine_captures_by_ref_ty, + ] = **obligation.predicate.args + else { + bug!(); + }; + + let predicate = ty::ProjectionPredicate { + projection_ty: ty::AliasTy::new( + selcx.tcx(), + obligation.predicate.def_id, + obligation.predicate.args, + ), + term: ty::CoroutineClosureSignature::tupled_upvars_by_closure_kind( + selcx.tcx(), + goal_kind_ty.expect_ty().to_opt_closure_kind().unwrap(), + tupled_inputs_ty.expect_ty(), + tupled_upvars_ty.expect_ty(), + coroutine_captures_by_ref_ty.expect_ty(), + borrow_region.expect_region(), + ) + .into(), + }; + + confirm_param_env_candidate(selcx, obligation, ty::Binder::dummy(predicate), false) + .with_addl_obligations(nested) +} + fn confirm_param_env_candidate<'cx, 'tcx>( selcx: &mut SelectionContext<'cx, 'tcx>, obligation: &ProjectionTyObligation<'tcx>, diff --git a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs index b354ebf111f..75aedd5cd22 100644 --- a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs +++ b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs @@ -117,9 +117,12 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> { self.assemble_iterator_candidates(obligation, &mut candidates); } else if lang_items.async_iterator_trait() == Some(def_id) { self.assemble_async_iterator_candidates(obligation, &mut candidates); + } else if lang_items.async_fn_kind_helper() == Some(def_id) { + self.assemble_async_fn_kind_helper_candidates(obligation, &mut candidates); } self.assemble_closure_candidates(obligation, &mut candidates); + self.assemble_async_closure_candidates(obligation, &mut candidates); self.assemble_fn_pointer_candidates(obligation, &mut candidates); self.assemble_candidates_from_impls(obligation, &mut candidates); self.assemble_candidates_from_object_ty(obligation, &mut candidates); @@ -335,6 +338,49 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> { } } + fn assemble_async_closure_candidates( + &mut self, + obligation: &PolyTraitObligation<'tcx>, + candidates: &mut SelectionCandidateSet<'tcx>, + ) { + let Some(goal_kind) = + self.tcx().async_fn_trait_kind_from_def_id(obligation.predicate.def_id()) + else { + return; + }; + + match *obligation.self_ty().skip_binder().kind() { + ty::CoroutineClosure(_, args) => { + if let Some(closure_kind) = + args.as_coroutine_closure().kind_ty().to_opt_closure_kind() + && !closure_kind.extends(goal_kind) + { + return; + } + candidates.vec.push(AsyncClosureCandidate); + } + ty::Infer(ty::TyVar(_)) => { + candidates.ambiguous = true; + } + _ => {} + } + } + + fn assemble_async_fn_kind_helper_candidates( + &mut self, + obligation: &PolyTraitObligation<'tcx>, + candidates: &mut SelectionCandidateSet<'tcx>, + ) { + if let Some(closure_kind) = obligation.self_ty().skip_binder().to_opt_closure_kind() + && let Some(goal_kind) = + obligation.predicate.skip_binder().trait_ref.args.type_at(1).to_opt_closure_kind() + { + if closure_kind.extends(goal_kind) { + candidates.vec.push(AsyncFnKindHelperCandidate); + } + } + } + /// Implements one of the `Fn()` family for a fn pointer. fn assemble_fn_pointer_candidates( &mut self, diff --git a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs index 74f388e53a3..79336a9d4e8 100644 --- a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs +++ b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs @@ -83,6 +83,13 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> { ImplSource::Builtin(BuiltinImplSource::Misc, vtable_closure) } + AsyncClosureCandidate => { + let vtable_closure = self.confirm_async_closure_candidate(obligation)?; + ImplSource::Builtin(BuiltinImplSource::Misc, vtable_closure) + } + + AsyncFnKindHelperCandidate => ImplSource::Builtin(BuiltinImplSource::Misc, vec![]), + CoroutineCandidate => { let vtable_coroutine = self.confirm_coroutine_candidate(obligation)?; ImplSource::Builtin(BuiltinImplSource::Misc, vtable_coroutine) @@ -869,6 +876,49 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> { Ok(nested) } + #[instrument(skip(self), level = "debug")] + fn confirm_async_closure_candidate( + &mut self, + obligation: &PolyTraitObligation<'tcx>, + ) -> Result>, SelectionError<'tcx>> { + // Okay to skip binder because the args on closure types never + // touch bound regions, they just capture the in-scope + // type/region parameters. + let self_ty = self.infcx.shallow_resolve(obligation.self_ty().skip_binder()); + let ty::CoroutineClosure(closure_def_id, args) = *self_ty.kind() else { + bug!("async closure candidate for non-coroutine-closure {:?}", obligation); + }; + + let trait_ref = args.as_coroutine_closure().coroutine_closure_sig().map_bound(|sig| { + ty::TraitRef::new( + self.tcx(), + obligation.predicate.def_id(), + [self_ty, sig.tupled_inputs_ty], + ) + }); + + let mut nested = self.confirm_poly_trait_refs(obligation, trait_ref)?; + + let goal_kind = + self.tcx().async_fn_trait_kind_from_def_id(obligation.predicate.def_id()).unwrap(); + nested.push(obligation.with( + self.tcx(), + ty::TraitRef::from_lang_item( + self.tcx(), + LangItem::AsyncFnKindHelper, + obligation.cause.span, + [ + args.as_coroutine_closure().kind_ty(), + Ty::from_closure_kind(self.tcx(), goal_kind), + ], + ), + )); + + debug!(?closure_def_id, ?trait_ref, ?nested, "confirm closure candidate obligations"); + + Ok(nested) + } + /// In the case of closure types and fn pointers, /// we currently treat the input type parameters on the trait as /// outputs. This means that when we have a match we have only diff --git a/compiler/rustc_trait_selection/src/traits/select/mod.rs b/compiler/rustc_trait_selection/src/traits/select/mod.rs index e79277b89c4..7f41c73b72f 100644 --- a/compiler/rustc_trait_selection/src/traits/select/mod.rs +++ b/compiler/rustc_trait_selection/src/traits/select/mod.rs @@ -1864,6 +1864,8 @@ impl<'tcx> SelectionContext<'_, 'tcx> { ImplCandidate(..) | AutoImplCandidate | ClosureCandidate { .. } + | AsyncClosureCandidate + | AsyncFnKindHelperCandidate | CoroutineCandidate | FutureCandidate | IteratorCandidate @@ -1894,6 +1896,8 @@ impl<'tcx> SelectionContext<'_, 'tcx> { ImplCandidate(_) | AutoImplCandidate | ClosureCandidate { .. } + | AsyncClosureCandidate + | AsyncFnKindHelperCandidate | CoroutineCandidate | FutureCandidate | IteratorCandidate @@ -1930,6 +1934,8 @@ impl<'tcx> SelectionContext<'_, 'tcx> { ImplCandidate(..) | AutoImplCandidate | ClosureCandidate { .. } + | AsyncClosureCandidate + | AsyncFnKindHelperCandidate | CoroutineCandidate | FutureCandidate | IteratorCandidate @@ -1946,6 +1952,8 @@ impl<'tcx> SelectionContext<'_, 'tcx> { ImplCandidate(..) | AutoImplCandidate | ClosureCandidate { .. } + | AsyncClosureCandidate + | AsyncFnKindHelperCandidate | CoroutineCandidate | FutureCandidate | IteratorCandidate @@ -2054,6 +2062,8 @@ impl<'tcx> SelectionContext<'_, 'tcx> { ( ImplCandidate(_) | ClosureCandidate { .. } + | AsyncClosureCandidate + | AsyncFnKindHelperCandidate | CoroutineCandidate | FutureCandidate | IteratorCandidate @@ -2066,6 +2076,8 @@ impl<'tcx> SelectionContext<'_, 'tcx> { | TraitAliasCandidate, ImplCandidate(_) | ClosureCandidate { .. } + | AsyncClosureCandidate + | AsyncFnKindHelperCandidate | CoroutineCandidate | FutureCandidate | IteratorCandidate diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs index af26616831f..8d1bdec6684 100644 --- a/compiler/rustc_ty_utils/src/abi.rs +++ b/compiler/rustc_ty_utils/src/abi.rs @@ -101,6 +101,42 @@ fn fn_sig_for_fn_abi<'tcx>( bound_vars, ) } + ty::CoroutineClosure(def_id, args) => { + let sig = args.as_coroutine_closure().coroutine_closure_sig(); + let bound_vars = tcx.mk_bound_variable_kinds_from_iter( + sig.bound_vars().iter().chain(iter::once(ty::BoundVariableKind::Region(ty::BrEnv))), + ); + let br = ty::BoundRegion { + var: ty::BoundVar::from_usize(bound_vars.len() - 1), + kind: ty::BoundRegionKind::BrEnv, + }; + let env_region = ty::Region::new_bound(tcx, ty::INNERMOST, br); + let env_ty = tcx.closure_env_ty( + Ty::new_coroutine_closure(tcx, def_id, args), + args.as_coroutine_closure().kind(), + env_region, + ); + + let sig = sig.skip_binder(); + ty::Binder::bind_with_vars( + tcx.mk_fn_sig( + iter::once(env_ty).chain([sig.tupled_inputs_ty]), + sig.to_coroutine_given_kind_and_upvars( + tcx, + args.as_coroutine_closure().parent_args(), + tcx.coroutine_for_closure(def_id), + args.as_coroutine_closure().kind(), + env_region, + args.as_coroutine_closure().tupled_upvars_ty(), + args.as_coroutine_closure().coroutine_captures_by_ref_ty(), + ), + sig.c_variadic, + sig.unsafety, + sig.abi, + ), + bound_vars, + ) + } ty::Coroutine(did, args) => { let coroutine_kind = tcx.coroutine_kind(did).unwrap(); let sig = args.as_coroutine().sig(); diff --git a/compiler/rustc_ty_utils/src/instance.rs b/compiler/rustc_ty_utils/src/instance.rs index a5a2c53c732..50d27ed4ced 100644 --- a/compiler/rustc_ty_utils/src/instance.rs +++ b/compiler/rustc_ty_utils/src/instance.rs @@ -38,6 +38,7 @@ fn resolve_instance<'tcx>( debug!(" => nontrivial drop glue"); match *ty.kind() { ty::Closure(..) + | ty::CoroutineClosure(..) | ty::Coroutine(..) | ty::Tuple(..) | ty::Adt(..) @@ -282,6 +283,16 @@ fn resolve_associated_item<'tcx>( tcx.item_name(trait_item_id) ), } + } else if tcx.async_fn_trait_kind_from_def_id(trait_ref.def_id).is_some() { + match *rcvr_args.type_at(0).kind() { + ty::CoroutineClosure(closure_def_id, args) => { + Some(Instance::new(closure_def_id, args)) + } + _ => bug!( + "no built-in definition for `{trait_ref}::{}` for non-lending-closure type", + tcx.item_name(trait_item_id) + ), + } } else { Instance::try_resolve_item_for_coroutine(tcx, trait_item_id, trait_id, rcvr_args) } diff --git a/library/core/src/ops/async_function.rs b/library/core/src/ops/async_function.rs index 965873f163e..b11d5643990 100644 --- a/library/core/src/ops/async_function.rs +++ b/library/core/src/ops/async_function.rs @@ -106,3 +106,11 @@ mod impls { } } } + +mod internal_implementation_detail { + // TODO: needs a detailed explanation + #[cfg_attr(not(bootstrap), lang = "async_fn_kind_helper")] + trait AsyncFnKindHelper { + type Assoc<'closure_env, Inputs, Upvars, BorrowedUpvarsAsFnPtr>; + } +} -- cgit 1.4.1-3-g733a5 From c98d6994a390044410c55c45195d330e4c8cd3d7 Mon Sep 17 00:00:00 2001 From: Michael Goulet Date: Thu, 25 Jan 2024 19:17:21 +0000 Subject: More comments, final tweaks --- compiler/rustc_ast_lowering/src/expr.rs | 3 ++ .../rustc_borrowck/src/type_check/input_output.rs | 8 +++- compiler/rustc_borrowck/src/universal_regions.rs | 8 ++++ compiler/rustc_hir_analysis/src/collect.rs | 11 +----- compiler/rustc_hir_typeck/src/callee.rs | 9 +++++ compiler/rustc_hir_typeck/src/closure.rs | 39 +++++++++++++------ compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs | 2 +- compiler/rustc_hir_typeck/src/upvar.rs | 10 ++++- compiler/rustc_middle/src/ty/flags.rs | 6 +-- compiler/rustc_middle/src/ty/instance.rs | 7 ++++ compiler/rustc_middle/src/ty/sty.rs | 44 +++++++++++++++++++--- compiler/rustc_mir_dataflow/src/value_analysis.rs | 3 ++ .../src/coroutine/by_move_body.rs | 13 ++++++- .../src/solve/assembly/structural_traits.rs | 24 +++++++++--- .../rustc_trait_selection/src/traits/project.rs | 3 ++ .../src/traits/select/candidate_assembly.rs | 4 ++ .../src/traits/select/confirmation.rs | 2 + compiler/rustc_trait_selection/src/traits/wf.rs | 2 +- compiler/rustc_ty_utils/src/abi.rs | 6 +++ compiler/rustc_ty_utils/src/instance.rs | 7 ++++ 20 files changed, 169 insertions(+), 42 deletions(-) (limited to 'compiler/rustc_mir_dataflow') diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs index 9990d526bf7..337dd55541a 100644 --- a/compiler/rustc_ast_lowering/src/expr.rs +++ b/compiler/rustc_ast_lowering/src/expr.rs @@ -1074,6 +1074,9 @@ impl<'hir> LoweringContext<'_, 'hir> { body, fn_decl_span: self.lower_span(fn_decl_span), fn_arg_span: Some(self.lower_span(fn_arg_span)), + // Lower this as a `CoroutineClosure`. That will ensure that HIR typeck + // knows that a `FnDecl` output type like `-> &str` actually means + // "coroutine that returns &str", rather than directly returning a `&str`. kind: hir::ClosureKind::CoroutineClosure(hir::CoroutineDesugaring::Async), constness: hir::Constness::NotConst, }); diff --git a/compiler/rustc_borrowck/src/type_check/input_output.rs b/compiler/rustc_borrowck/src/type_check/input_output.rs index 46fd0ce2b39..a5a7ce4ea3e 100644 --- a/compiler/rustc_borrowck/src/type_check/input_output.rs +++ b/compiler/rustc_borrowck/src/type_check/input_output.rs @@ -61,6 +61,8 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> { )), "this needs to be modified if we're lowering non-async closures" ); + // Make sure to use the args from `DefiningTy` so the right NLL region vids are prepopulated + // into the type. let args = args.as_coroutine_closure(); let tupled_upvars_ty = ty::CoroutineClosureSignature::tupled_upvars_by_closure_kind( self.tcx(), @@ -87,11 +89,13 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> { ty::CoroutineArgsParts { parent_args: args.parent_args(), kind_ty: Ty::from_closure_kind(self.tcx(), args.kind()), + return_ty: user_provided_sig.output(), + tupled_upvars_ty, + // For async closures, none of these can be annotated, so just fill + // them with fresh ty vars. resume_ty: next_ty_var(), yield_ty: next_ty_var(), witness: next_ty_var(), - return_ty: user_provided_sig.output(), - tupled_upvars_ty: tupled_upvars_ty, }, ) .args, diff --git a/compiler/rustc_borrowck/src/universal_regions.rs b/compiler/rustc_borrowck/src/universal_regions.rs index d696e624823..111eaaf60f7 100644 --- a/compiler/rustc_borrowck/src/universal_regions.rs +++ b/compiler/rustc_borrowck/src/universal_regions.rs @@ -720,6 +720,14 @@ impl<'cx, 'tcx> UniversalRegionsBuilder<'cx, 'tcx> { ty::Binder::dummy(inputs_and_output) } + // Construct the signature of the CoroutineClosure for the purposes of borrowck. + // This is pretty straightforward -- we: + // 1. first grab the `coroutine_closure_sig`, + // 2. compute the self type (`&`/`&mut`/no borrow), + // 3. flatten the tupled_input_tys, + // 4. construct the correct generator type to return with + // `CoroutineClosureSignature::to_coroutine_given_kind_and_upvars`. + // Then we wrap it all up into a list of inputs and output. DefiningTy::CoroutineClosure(def_id, args) => { assert_eq!(self.mir_def.to_def_id(), def_id); let closure_sig = args.as_coroutine_closure().coroutine_closure_sig(); diff --git a/compiler/rustc_hir_analysis/src/collect.rs b/compiler/rustc_hir_analysis/src/collect.rs index 8fdd3cfe3cf..fbcebb7c87c 100644 --- a/compiler/rustc_hir_analysis/src/collect.rs +++ b/compiler/rustc_hir_analysis/src/collect.rs @@ -1533,15 +1533,8 @@ fn coroutine_kind(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option, def_id: LocalDefId) -> DefId { - let Node::Expr(&hir::Expr { - kind: - hir::ExprKind::Closure(&rustc_hir::Closure { - kind: hir::ClosureKind::CoroutineClosure(_), - body, - .. - }), - .. - }) = tcx.hir_node_by_def_id(def_id) + let &rustc_hir::Closure { kind: hir::ClosureKind::CoroutineClosure(_), body, .. } = + tcx.hir_node_by_def_id(def_id).expect_closure() else { bug!() }; diff --git a/compiler/rustc_hir_typeck/src/callee.rs b/compiler/rustc_hir_typeck/src/callee.rs index 730a475f630..fbe6f454dbc 100644 --- a/compiler/rustc_hir_typeck/src/callee.rs +++ b/compiler/rustc_hir_typeck/src/callee.rs @@ -166,6 +166,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { return Some(CallStep::DeferredClosure(def_id, closure_sig)); } + // When calling a `CoroutineClosure` that is local to the body, we will + // not know what its `closure_kind` is yet. Instead, just fill in the + // signature with an infer var for the `tupled_upvars_ty` of the coroutine, + // and record a deferred call resolution which will constrain that var + // as part of `AsyncFn*` trait confirmation. ty::CoroutineClosure(def_id, args) if self.closure_kind(adjusted_ty).is_none() => { let def_id = def_id.expect_local(); let closure_args = args.as_coroutine_closure(); @@ -183,6 +188,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { 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(), self.tcx.coroutine_for_closure(def_id), tupled_upvars_ty, diff --git a/compiler/rustc_hir_typeck/src/closure.rs b/compiler/rustc_hir_typeck/src/closure.rs index 014293c1f83..a985fa201d0 100644 --- a/compiler/rustc_hir_typeck/src/closure.rs +++ b/compiler/rustc_hir_typeck/src/closure.rs @@ -175,6 +175,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { interior, )); + // Coroutines that come from coroutine closures have not yet determined + // their kind ty, so make a fresh infer var which will be constrained + // later during upvar analysis. Regular coroutines always have the kind + // ty of `().` let kind_ty = match kind { hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure) => self .next_ty_var(TypeVariableOrigin { @@ -203,6 +207,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { ) } hir::ClosureKind::CoroutineClosure(kind) => { + // async closures always return the type ascribed after the `->` (if present), + // and yield `()`. let (bound_return_ty, bound_yield_ty) = match kind { hir::CoroutineDesugaring::Async => { (bound_sig.skip_binder().output(), tcx.types.unit) @@ -211,6 +217,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { todo!("`gen` and `async gen` closures not supported yet") } }; + // Compute all of the variables that will be used to populate the coroutine. let resume_ty = self.next_ty_var(TypeVariableOrigin { kind: TypeVariableOriginKind::ClosureSynthetic, span: expr_span, @@ -258,20 +265,28 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { kind: TypeVariableOriginKind::ClosureSynthetic, span: expr_span, }); + + // We need to turn the liberated signature that we got from HIR, which + // looks something like `|Args...| -> T`, into a signature that is suitable + // for type checking the inner body of the closure, which always returns a + // coroutine. To do so, we use the `CoroutineClosureSignature` to compute + // the coroutine type, filling in the tupled_upvars_ty and kind_ty with infer + // vars which will get constrained during upvar analysis. + let coroutine_output_ty = tcx.liberate_late_bound_regions( + expr_def_id.to_def_id(), + closure_args.coroutine_closure_sig().map_bound(|sig| { + sig.to_coroutine( + tcx, + parent_args, + closure_kind_ty, + tcx.coroutine_for_closure(expr_def_id), + coroutine_upvars_ty, + ) + }), + ); liberated_sig = tcx.mk_fn_sig( liberated_sig.inputs().iter().copied(), - tcx.liberate_late_bound_regions( - expr_def_id.to_def_id(), - closure_args.coroutine_closure_sig().map_bound(|sig| { - sig.to_coroutine( - tcx, - parent_args, - closure_kind_ty, - tcx.coroutine_for_closure(expr_def_id), - coroutine_upvars_ty, - ) - }), - ), + coroutine_output_ty, liberated_sig.c_variadic, liberated_sig.unsafety, liberated_sig.abi, diff --git a/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs b/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs index 23fb7eba656..d30c7a4fb38 100644 --- a/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs +++ b/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs @@ -335,7 +335,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { continue; } - // For this check, we do *not* want to treat async coroutine-closures (async blocks) + // For this check, we do *not* want to treat async coroutine closures (async blocks) // as proper closures. Doing so would regress type inference when feeding // the return value of an argument-position async block to an argument-position // closure wrapped in a block. diff --git a/compiler/rustc_hir_typeck/src/upvar.rs b/compiler/rustc_hir_typeck/src/upvar.rs index e8bdf283c4f..c4773a88521 100644 --- a/compiler/rustc_hir_typeck/src/upvar.rs +++ b/compiler/rustc_hir_typeck/src/upvar.rs @@ -336,6 +336,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { } } + // For coroutine-closures, we additionally must compute the + // `coroutine_captures_by_ref_ty` type, which is used to generate the by-ref + // version of the coroutine-closure's output coroutine. if let UpvarArgs::CoroutineClosure(args) = args { let closure_env_region: ty::Region<'_> = ty::Region::new_bound( self.tcx, @@ -353,7 +356,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { self.tcx.coroutine_for_closure(closure_def_id).expect_local(), ) // Skip the captures that are just moving the closure's args - // into the coroutine. These are always by move. + // into the coroutine. These are always by move, and we append + // those later in the `CoroutineClosureSignature` helper functions. .skip( args.as_coroutine_closure() .coroutine_closure_sig() @@ -365,6 +369,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { .map(|captured_place| { let upvar_ty = captured_place.place.ty(); let capture = captured_place.info.capture_kind; + // Not all upvars are captured by ref, so use + // `apply_capture_kind_on_capture_ty` to ensure that we + // compute the right captured type. apply_capture_kind_on_capture_ty( self.tcx, upvar_ty, @@ -394,6 +401,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> { coroutine_captures_by_ref_ty, ); + // Additionally, we can now constrain the coroutine's kind type. let ty::Coroutine(_, coroutine_args) = *self.typeck_results.borrow().expr_ty(body.value).kind() else { diff --git a/compiler/rustc_middle/src/ty/flags.rs b/compiler/rustc_middle/src/ty/flags.rs index deb6e73ea71..0f4b5fe228c 100644 --- a/compiler/rustc_middle/src/ty/flags.rs +++ b/compiler/rustc_middle/src/ty/flags.rs @@ -145,11 +145,11 @@ impl FlagComputation { self.flags -= TypeFlags::STILL_FURTHER_SPECIALIZABLE; } - self.add_ty(args.signature_parts_ty()); - self.add_ty(args.coroutine_witness_ty()); - self.add_ty(args.coroutine_captures_by_ref_ty()); self.add_ty(args.kind_ty()); + self.add_ty(args.signature_parts_ty()); self.add_ty(args.tupled_upvars_ty()); + self.add_ty(args.coroutine_captures_by_ref_ty()); + self.add_ty(args.coroutine_witness_ty()); } &ty::Bound(debruijn, _) => { diff --git a/compiler/rustc_middle/src/ty/instance.rs b/compiler/rustc_middle/src/ty/instance.rs index 2c80dd02145..1c9415ef3b0 100644 --- a/compiler/rustc_middle/src/ty/instance.rs +++ b/compiler/rustc_middle/src/ty/instance.rs @@ -646,9 +646,16 @@ impl<'tcx> Instance<'tcx> { bug!() }; + // If the closure's kind ty disagrees with the identity closure's kind ty, + // then this must be a coroutine generated by one of the `ConstructCoroutineInClosureShim`s. if args.as_coroutine().kind_ty() == id_args.as_coroutine().kind_ty() { Some(Instance { def: ty::InstanceDef::Item(coroutine_def_id), args }) } else { + assert_eq!( + args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(), + ty::ClosureKind::FnOnce, + "FIXME(async_closures): Generate a by-mut body here." + ); Some(Instance { def: ty::InstanceDef::CoroutineByMoveShim { coroutine_def_id }, args, diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs index 3bbebaddbdd..927924452f9 100644 --- a/compiler/rustc_middle/src/ty/sty.rs +++ b/compiler/rustc_middle/src/ty/sty.rs @@ -275,6 +275,10 @@ pub struct CoroutineClosureArgs<'tcx> { pub args: GenericArgsRef<'tcx>, } +/// See docs for explanation of how each argument is used. +/// +/// See [`CoroutineClosureSignature`] for how these arguments are put together +/// to make a callable [`FnSig`] suitable for typeck and borrowck. pub struct CoroutineClosureArgsParts<'tcx> { /// This is the args of the typeck root. pub parent_args: &'tcx [GenericArg<'tcx>], @@ -409,17 +413,32 @@ pub struct CoroutineClosureSignature<'tcx> { pub resume_ty: Ty<'tcx>, pub yield_ty: Ty<'tcx>, pub return_ty: Ty<'tcx>, + + // Like the `fn_sig_as_fn_ptr_ty` of a regular closure, these types + // never actually differ. But we save them rather than recreating them + // from scratch just for good measure. + /// Always false pub c_variadic: bool, + /// Always [`hir::Unsafety::Normal`] pub unsafety: hir::Unsafety, + /// Always [`abi::Abi::RustCall`] pub abi: abi::Abi, } impl<'tcx> CoroutineClosureSignature<'tcx> { + /// Construct a coroutine from the closure signature. Since a coroutine signature + /// is agnostic to the type of generator that is returned (by-ref/by-move), + /// the caller must specify what "flavor" of generator that they'd like to + /// create. Additionally, they must manually compute the upvars of the closure. + /// + /// This helper is not really meant to be used directly except for early on + /// during typeck, when we want to put inference vars into the kind and upvars tys. + /// When the kind and upvars are known, use the other helper functions. pub fn to_coroutine( self, tcx: TyCtxt<'tcx>, parent_args: &'tcx [GenericArg<'tcx>], - kind_ty: Ty<'tcx>, + coroutine_kind_ty: Ty<'tcx>, coroutine_def_id: DefId, tupled_upvars_ty: Ty<'tcx>, ) -> Ty<'tcx> { @@ -427,7 +446,7 @@ impl<'tcx> CoroutineClosureSignature<'tcx> { tcx, ty::CoroutineArgsParts { parent_args, - kind_ty, + kind_ty: coroutine_kind_ty, resume_ty: self.resume_ty, yield_ty: self.yield_ty, return_ty: self.return_ty, @@ -439,19 +458,24 @@ impl<'tcx> CoroutineClosureSignature<'tcx> { Ty::new_coroutine(tcx, coroutine_def_id, coroutine_args.args) } + /// Given known upvars and a [`ClosureKind`](ty::ClosureKind), compute the coroutine + /// returned by that corresponding async fn trait. + /// + /// This function expects the upvars to have been computed already, and doesn't check + /// that the `ClosureKind` is actually supported by the coroutine-closure. pub fn to_coroutine_given_kind_and_upvars( self, tcx: TyCtxt<'tcx>, parent_args: &'tcx [GenericArg<'tcx>], coroutine_def_id: DefId, - closure_kind: ty::ClosureKind, + goal_kind: ty::ClosureKind, env_region: ty::Region<'tcx>, closure_tupled_upvars_ty: Ty<'tcx>, coroutine_captures_by_ref_ty: Ty<'tcx>, ) -> Ty<'tcx> { let tupled_upvars_ty = Self::tupled_upvars_by_closure_kind( tcx, - closure_kind, + goal_kind, self.tupled_inputs_ty, closure_tupled_upvars_ty, coroutine_captures_by_ref_ty, @@ -461,13 +485,21 @@ impl<'tcx> CoroutineClosureSignature<'tcx> { self.to_coroutine( tcx, parent_args, - Ty::from_closure_kind(tcx, closure_kind), + Ty::from_closure_kind(tcx, goal_kind), coroutine_def_id, tupled_upvars_ty, ) } - /// Given a closure kind, compute the tupled upvars that the given coroutine would return. + /// Compute the tupled upvars that a coroutine-closure's output coroutine + /// would return for the given `ClosureKind`. + /// + /// When `ClosureKind` is `FnMut`/`Fn`, then this will use the "captures by ref" + /// to return a set of upvars which are borrowed with the given `env_region`. + /// + /// This ensures that the `AsyncFn::call` will return a coroutine whose upvars' + /// lifetimes are related to the lifetime of the borrow on the closure made for + /// the call. This allows borrowck to enforce the self-borrows correctly. pub fn tupled_upvars_by_closure_kind( tcx: TyCtxt<'tcx>, kind: ty::ClosureKind, diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs index a85e53ff241..2b2af6ee7da 100644 --- a/compiler/rustc_mir_dataflow/src/value_analysis.rs +++ b/compiler/rustc_mir_dataflow/src/value_analysis.rs @@ -1149,6 +1149,9 @@ pub fn iter_fields<'tcx>( ty::Closure(_, args) => { iter_fields(args.as_closure().tupled_upvars_ty(), tcx, param_env, f); } + ty::Coroutine(_, args) => { + iter_fields(args.as_coroutine().tupled_upvars_ty(), tcx, param_env, f); + } ty::CoroutineClosure(_, args) => { iter_fields(args.as_coroutine_closure().tupled_upvars_ty(), tcx, param_env, f); } diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs index 4e3e70bdafe..1cc0a5026d1 100644 --- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs +++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs @@ -1,3 +1,8 @@ +//! 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; @@ -87,11 +92,16 @@ impl<'tcx> MutVisitor<'tcx> for MakeByMoveBody<'tcx> { && 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, ty.builtin_deref(true).unwrap().ty)] + [mir::ProjectionElem::Field(idx, peeled_ty)] .into_iter() .chain(end.iter().copied()), ), @@ -101,6 +111,7 @@ impl<'tcx> MutVisitor<'tcx> for MakeByMoveBody<'tcx> { } 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_trait_selection/src/solve/assembly/structural_traits.rs b/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs index dbf3e4876a9..d02578c4846 100644 --- a/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs +++ b/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs @@ -309,14 +309,19 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_callable<'tcx>( } // Returns a binder of the tupled inputs types, output type, and coroutine type -// from a builtin async closure type. +// from a builtin coroutine-closure type. If we don't yet know the closure kind of +// the coroutine-closure, emit an additional trait predicate for `AsyncFnKindHelper` +// which enforces the closure is actually callable with the given trait. When we +// know the kind already, we can short-circuit this check. pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tcx>( tcx: TyCtxt<'tcx>, self_ty: Ty<'tcx>, goal_kind: ty::ClosureKind, env_region: ty::Region<'tcx>, -) -> Result<(ty::Binder<'tcx, (Ty<'tcx>, Ty<'tcx>, Ty<'tcx>)>, Vec>), NoSolution> -{ +) -> Result< + (ty::Binder<'tcx, (Ty<'tcx>, Ty<'tcx>, Ty<'tcx>)>, Option>), + NoSolution, +> { match *self_ty.kind() { ty::CoroutineClosure(def_id, args) => { let args = args.as_coroutine_closure(); @@ -339,7 +344,7 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tc ); (sig.tupled_inputs_ty, sig.return_ty, coroutine_ty) }), - vec![], + None, )) } else { let async_fn_kind_trait_def_id = @@ -350,6 +355,13 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tc .next() .unwrap() .def_id; + // When we don't know the closure kind (and therefore also the closure's upvars, + // which are computed at the same time), we must delay the computation of the + // generator's upvars. We do this using the `AsyncFnKindHelper`, which as a trait + // goal functions similarly to the old `ClosureKind` predicate, and ensures that + // the goal kind <= the closure kind. As a projection `AsyncFnKindHelper::Upvars` + // will project to the right upvars for the generator, appending the inputs and + // coroutine upvars respecting the closure kind. Ok(( args.coroutine_closure_sig().map_bound(|sig| { let tupled_upvars_ty = Ty::new_projection( @@ -373,14 +385,14 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tc ); (sig.tupled_inputs_ty, sig.return_ty, coroutine_ty) }), - vec![ + Some( ty::TraitRef::new( tcx, async_fn_kind_trait_def_id, [kind_ty, Ty::from_closure_kind(tcx, goal_kind)], ) .to_predicate(tcx), - ], + ), )) } } diff --git a/compiler/rustc_trait_selection/src/traits/project.rs b/compiler/rustc_trait_selection/src/traits/project.rs index a8d6b9812be..955c81eee6b 100644 --- a/compiler/rustc_trait_selection/src/traits/project.rs +++ b/compiler/rustc_trait_selection/src/traits/project.rs @@ -2491,6 +2491,9 @@ fn confirm_async_closure_candidate<'cx, 'tcx>( // since all this does is make the solver do more work. // // The code duplication due to the different length args is kind of weird, too. + // + // See the logic in `structural_traits` in the new solver to understand a bit + // more clearly how this *should* look. let poly_cache_entry = args.coroutine_closure_sig().map_bound(|sig| { let (projection_ty, term) = match tcx.item_name(obligation.predicate.def_id) { sym::CallOnceFuture => { diff --git a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs index dda68fd4244..2258e796103 100644 --- a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs +++ b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs @@ -121,9 +121,11 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> { self.assemble_async_fn_kind_helper_candidates(obligation, &mut candidates); } + // FIXME: Put these into `else if` blocks above, since they're built-in. self.assemble_closure_candidates(obligation, &mut candidates); self.assemble_async_closure_candidates(obligation, &mut candidates); self.assemble_fn_pointer_candidates(obligation, &mut candidates); + self.assemble_candidates_from_impls(obligation, &mut candidates); self.assemble_candidates_from_object_ty(obligation, &mut candidates); } @@ -382,6 +384,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> { return; } + // Check that the self kind extends the goal kind. If it does, + // then there's nothing else to check. if let Some(closure_kind) = self_ty.to_opt_closure_kind() && let Some(goal_kind) = target_kind_ty.to_opt_closure_kind() { diff --git a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs index 79336a9d4e8..c9d06b0f675 100644 --- a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs +++ b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs @@ -88,6 +88,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> { ImplSource::Builtin(BuiltinImplSource::Misc, vtable_closure) } + // No nested obligations or confirmation process. The checks that we do in + // candidate assembly are sufficient. AsyncFnKindHelperCandidate => ImplSource::Builtin(BuiltinImplSource::Misc, vec![]), CoroutineCandidate => { diff --git a/compiler/rustc_trait_selection/src/traits/wf.rs b/compiler/rustc_trait_selection/src/traits/wf.rs index 88584a61b13..b4f13ee95a6 100644 --- a/compiler/rustc_trait_selection/src/traits/wf.rs +++ b/compiler/rustc_trait_selection/src/traits/wf.rs @@ -728,7 +728,7 @@ impl<'a, 'tcx> WfPredicates<'a, 'tcx> { } ty::CoroutineClosure(did, args) => { - // See the above comments. + // See the above comments. The same apply to coroutine-closures. walker.skip_current_subtree(); self.compute(args.as_coroutine_closure().tupled_upvars_ty().into()); let obligations = self.nominal_obligations(did, args); diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs index 3ea48ee824d..e023283a709 100644 --- a/compiler/rustc_ty_utils/src/abi.rs +++ b/compiler/rustc_ty_utils/src/abi.rs @@ -112,6 +112,10 @@ fn fn_sig_for_fn_abi<'tcx>( }; let env_region = ty::Region::new_bound(tcx, ty::INNERMOST, br); + // When this `CoroutineClosure` comes from a `ConstructCoroutineInClosureShim`, + // make sure we respect the `target_kind` in that shim. + // FIXME(async_closures): This shouldn't be needed, and we should be populating + // a separate def-id for these bodies. let mut kind = args.as_coroutine_closure().kind(); if let InstanceDef::ConstructCoroutineInClosureShim { target_kind, .. } = instance.def { kind = target_kind; @@ -141,6 +145,8 @@ fn fn_sig_for_fn_abi<'tcx>( ) } ty::Coroutine(did, args) => { + // FIXME(async_closures): This isn't right for `CoroutineByMoveShim`. + let coroutine_kind = tcx.coroutine_kind(did).unwrap(); let sig = args.as_coroutine().sig(); diff --git a/compiler/rustc_ty_utils/src/instance.rs b/compiler/rustc_ty_utils/src/instance.rs index fc224d242db..9faad10dd14 100644 --- a/compiler/rustc_ty_utils/src/instance.rs +++ b/compiler/rustc_ty_utils/src/instance.rs @@ -287,6 +287,13 @@ fn resolve_associated_item<'tcx>( { match *rcvr_args.type_at(0).kind() { ty::CoroutineClosure(coroutine_closure_def_id, args) => { + // If we're computing `AsyncFnOnce`/`AsyncFnMut` for a by-ref closure, + // or `AsyncFnOnce` for a by-mut closure, then construct a new body that + // has the right return types. + // + // Specifically, `AsyncFnMut` for a by-ref coroutine-closure just needs + // to have its input and output types fixed (`&mut self` and returning + // `i16` coroutine kind). if target_kind > args.as_coroutine_closure().kind() { Some(Instance { def: ty::InstanceDef::ConstructCoroutineInClosureShim { -- cgit 1.4.1-3-g733a5