Skip to content

Commit 4049c3b

Browse files
ChayimFriedman2ChayimFriedman2
committed
Support AsyncFnX traits
Only in calls, because to support them in bounds we need support from Chalk. However we don't yet report error from bounds anyway, so this is less severe. The returned future is shown in its name within inlay hints instead of as a nicer `impl Future`, but that can wait for another PR.
1 parent e6276c8 commit 4049c3b

File tree

9 files changed

+291
-68
lines changed

9 files changed

+291
-68
lines changed

‎crates/hir-def/src/lang_item.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,9 @@ language_item_table! {
376376
Fn, sym::fn_, fn_trait, Target::Trait, GenericRequirement::Exact(1);
377377
FnMut, sym::fn_mut, fn_mut_trait, Target::Trait, GenericRequirement::Exact(1);
378378
FnOnce, sym::fn_once, fn_once_trait, Target::Trait, GenericRequirement::Exact(1);
379+
AsyncFn, sym::async_fn, async_fn_trait, Target::Trait, GenericRequirement::Exact(1);
380+
AsyncFnMut, sym::async_fn_mut, async_fn_mut_trait, Target::Trait, GenericRequirement::Exact(1);
381+
AsyncFnOnce, sym::async_fn_once, async_fn_once_trait, Target::Trait, GenericRequirement::Exact(1);
379382

380383
FnOnceOutput, sym::fn_once_output, fn_once_output, Target::AssocTy, GenericRequirement::None;
381384

‎crates/hir-ty/src/infer/expr.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1287,8 +1287,8 @@ impl InferenceContext<'_> {
12871287
tgt_expr: ExprId,
12881288
) {
12891289
match fn_x {
1290-
FnTrait::FnOnce => (),
1291-
FnTrait::FnMut => {
1290+
FnTrait::FnOnce | FnTrait::AsyncFnOnce => (),
1291+
FnTrait::FnMut | FnTrait::AsyncFnMut => {
12921292
if let TyKind::Ref(Mutability::Mut, lt, inner) = derefed_callee.kind(Interner) {
12931293
if adjustments
12941294
.last()
@@ -1312,7 +1312,7 @@ impl InferenceContext<'_> {
13121312
));
13131313
}
13141314
}
1315-
FnTrait::Fn => {
1315+
FnTrait::Fn | FnTrait::AsyncFn => {
13161316
if !matches!(derefed_callee.kind(Interner), TyKind::Ref(Mutability::Not, _, _)) {
13171317
adjustments.push(Adjustment::borrow(
13181318
Mutability::Not,

‎crates/hir-ty/src/infer/unify.rs

Lines changed: 60 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -794,69 +794,75 @@ impl<'a> InferenceTable<'a> {
794794
ty: &Ty,
795795
num_args: usize,
796796
) -> Option<(FnTrait, Vec<Ty>, Ty)> {
797-
let krate = self.trait_env.krate;
798-
let fn_once_trait = FnTrait::FnOnce.get_id(self.db, krate)?;
799-
let trait_data = self.db.trait_data(fn_once_trait);
800-
let output_assoc_type =
801-
trait_data.associated_type_by_name(&Name::new_symbol_root(sym::Output.clone()))?;
802-
803-
let mut arg_tys = Vec::with_capacity(num_args);
804-
let arg_ty = TyBuilder::tuple(num_args)
805-
.fill(|it| {
806-
let arg = match it {
807-
ParamKind::Type => self.new_type_var(),
808-
ParamKind::Lifetime => unreachable!("Tuple with lifetime parameter"),
809-
ParamKind::Const(_) => unreachable!("Tuple with const parameter"),
810-
};
811-
arg_tys.push(arg.clone());
812-
arg.cast(Interner)
813-
})
814-
.build();
815-
816-
let b = TyBuilder::trait_ref(self.db, fn_once_trait);
817-
if b.remaining() != 2 {
818-
return None;
819-
}
820-
let mut trait_ref = b.push(ty.clone()).push(arg_ty).build();
797+
for (fn_trait_name, output_assoc_name, subtraits) in [
798+
(FnTrait::FnOnce, sym::Output.clone(), &[FnTrait::Fn, FnTrait::FnMut][..]),
799+
(FnTrait::AsyncFnMut, sym::CallRefFuture.clone(), &[FnTrait::AsyncFn]),
800+
(FnTrait::AsyncFnOnce, sym::CallOnceFuture.clone(), &[]),
801+
] {
802+
let krate = self.trait_env.krate;
803+
let fn_trait = fn_trait_name.get_id(self.db, krate)?;
804+
let trait_data = self.db.trait_data(fn_trait);
805+
let output_assoc_type =
806+
trait_data.associated_type_by_name(&Name::new_symbol_root(output_assoc_name))?;
807+
808+
let mut arg_tys = Vec::with_capacity(num_args);
809+
let arg_ty = TyBuilder::tuple(num_args)
810+
.fill(|it| {
811+
let arg = match it {
812+
ParamKind::Type => self.new_type_var(),
813+
ParamKind::Lifetime => unreachable!("Tuple with lifetime parameter"),
814+
ParamKind::Const(_) => unreachable!("Tuple with const parameter"),
815+
};
816+
arg_tys.push(arg.clone());
817+
arg.cast(Interner)
818+
})
819+
.build();
820+
821+
let b = TyBuilder::trait_ref(self.db, fn_trait);
822+
if b.remaining() != 2 {
823+
return None;
824+
}
825+
let mut trait_ref = b.push(ty.clone()).push(arg_ty).build();
821826

822-
let projection = {
823-
TyBuilder::assoc_type_projection(
827+
let projection = TyBuilder::assoc_type_projection(
824828
self.db,
825829
output_assoc_type,
826830
Some(trait_ref.substitution.clone()),
827831
)
828-
.build()
829-
};
832+
.fill_with_unknown()
833+
.build();
830834

831-
let trait_env = self.trait_env.env.clone();
832-
let obligation = InEnvironment {
833-
goal: trait_ref.clone().cast(Interner),
834-
environment: trait_env.clone(),
835-
};
836-
let canonical = self.canonicalize(obligation.clone());
837-
if self.db.trait_solve(krate, self.trait_env.block, canonical.cast(Interner)).is_some() {
838-
self.register_obligation(obligation.goal);
839-
let return_ty = self.normalize_projection_ty(projection);
840-
for fn_x in [FnTrait::Fn, FnTrait::FnMut, FnTrait::FnOnce] {
841-
let fn_x_trait = fn_x.get_id(self.db, krate)?;
842-
trait_ref.trait_id = to_chalk_trait_id(fn_x_trait);
843-
let obligation: chalk_ir::InEnvironment<chalk_ir::Goal<Interner>> = InEnvironment {
844-
goal: trait_ref.clone().cast(Interner),
845-
environment: trait_env.clone(),
846-
};
847-
let canonical = self.canonicalize(obligation.clone());
848-
if self
849-
.db
850-
.trait_solve(krate, self.trait_env.block, canonical.cast(Interner))
851-
.is_some()
852-
{
853-
return Some((fn_x, arg_tys, return_ty));
835+
let trait_env = self.trait_env.env.clone();
836+
let obligation = InEnvironment {
837+
goal: trait_ref.clone().cast(Interner),
838+
environment: trait_env.clone(),
839+
};
840+
let canonical = self.canonicalize(obligation.clone());
841+
if self.db.trait_solve(krate, self.trait_env.block, canonical.cast(Interner)).is_some()
842+
{
843+
self.register_obligation(obligation.goal);
844+
let return_ty = self.normalize_projection_ty(projection);
845+
for &fn_x in subtraits {
846+
let fn_x_trait = fn_x.get_id(self.db, krate)?;
847+
trait_ref.trait_id = to_chalk_trait_id(fn_x_trait);
848+
let obligation: chalk_ir::InEnvironment<chalk_ir::Goal<Interner>> =
849+
InEnvironment {
850+
goal: trait_ref.clone().cast(Interner),
851+
environment: trait_env.clone(),
852+
};
853+
let canonical = self.canonicalize(obligation.clone());
854+
if self
855+
.db
856+
.trait_solve(krate, self.trait_env.block, canonical.cast(Interner))
857+
.is_some()
858+
{
859+
return Some((fn_x, arg_tys, return_ty));
860+
}
854861
}
862+
return Some((fn_trait_name, arg_tys, return_ty));
855863
}
856-
unreachable!("It should at least implement FnOnce at this point");
857-
} else {
858-
None
859864
}
865+
None
860866
}
861867

862868
pub(super) fn insert_type_vars<T>(&mut self, ty: T) -> T

‎crates/hir-ty/src/mir/lower.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2023,11 +2023,11 @@ pub fn mir_body_for_closure_query(
20232023
ctx.result.locals.alloc(Local { ty: infer[*root].clone() });
20242024
let closure_local = ctx.result.locals.alloc(Local {
20252025
ty: match kind {
2026-
FnTrait::FnOnce => infer[expr].clone(),
2027-
FnTrait::FnMut => {
2026+
FnTrait::FnOnce | FnTrait::AsyncFnOnce => infer[expr].clone(),
2027+
FnTrait::FnMut | FnTrait::AsyncFnMut => {
20282028
TyKind::Ref(Mutability::Mut, error_lifetime(), infer[expr].clone()).intern(Interner)
20292029
}
2030-
FnTrait::Fn => {
2030+
FnTrait::Fn | FnTrait::AsyncFn => {
20312031
TyKind::Ref(Mutability::Not, error_lifetime(), infer[expr].clone()).intern(Interner)
20322032
}
20332033
},
@@ -2055,8 +2055,10 @@ pub fn mir_body_for_closure_query(
20552055
let mut err = None;
20562056
let closure_local = ctx.result.locals.iter().nth(1).unwrap().0;
20572057
let closure_projection = match kind {
2058-
FnTrait::FnOnce => vec![],
2059-
FnTrait::FnMut | FnTrait::Fn => vec![ProjectionElem::Deref],
2058+
FnTrait::FnOnce | FnTrait::AsyncFnOnce => vec![],
2059+
FnTrait::FnMut | FnTrait::Fn | FnTrait::AsyncFnMut | FnTrait::AsyncFn => {
2060+
vec![ProjectionElem::Deref]
2061+
}
20602062
};
20612063
ctx.result.walk_places(|p, store| {
20622064
if let Some(it) = upvar_map.get(&p.local) {

‎crates/hir-ty/src/tests/traits.rs

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4811,3 +4811,53 @@ fn bar(v: *const ()) {
48114811
"#]],
48124812
);
48134813
}
4814+
4815+
#[test]
4816+
fn async_fn_traits() {
4817+
check_infer(
4818+
r#"
4819+
//- minicore: async_fn
4820+
async fn foo<T: AsyncFn(u32) -> i32>(a: T) {
4821+
let fut1 = a(0);
4822+
fut1.await;
4823+
}
4824+
async fn bar<T: AsyncFnMut(u32) -> i32>(mut b: T) {
4825+
let fut2 = b(0);
4826+
fut2.await;
4827+
}
4828+
async fn baz<T: AsyncFnOnce(u32) -> i32>(c: T) {
4829+
let fut3 = c(0);
4830+
fut3.await;
4831+
}
4832+
"#,
4833+
expect![[r#"
4834+
37..38 'a': T
4835+
43..83 '{ ...ait; }': ()
4836+
43..83 '{ ...ait; }': impl Future<Output = ()>
4837+
53..57 'fut1': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4838+
60..61 'a': T
4839+
60..64 'a(0)': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4840+
62..63 '0': u32
4841+
70..74 'fut1': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4842+
70..80 'fut1.await': i32
4843+
124..129 'mut b': T
4844+
134..174 '{ ...ait; }': ()
4845+
134..174 '{ ...ait; }': impl Future<Output = ()>
4846+
144..148 'fut2': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4847+
151..152 'b': T
4848+
151..155 'b(0)': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4849+
153..154 '0': u32
4850+
161..165 'fut2': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4851+
161..171 'fut2.await': i32
4852+
216..217 'c': T
4853+
222..262 '{ ...ait; }': ()
4854+
222..262 '{ ...ait; }': impl Future<Output = ()>
4855+
232..236 'fut3': AsyncFnOnce::CallOnceFuture<T, (u32,)>
4856+
239..240 'c': T
4857+
239..243 'c(0)': AsyncFnOnce::CallOnceFuture<T, (u32,)>
4858+
241..242 '0': u32
4859+
249..253 'fut3': AsyncFnOnce::CallOnceFuture<T, (u32,)>
4860+
249..259 'fut3.await': i32
4861+
"#]],
4862+
);
4863+
}

‎crates/hir-ty/src/traits.rs

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,10 @@ pub enum FnTrait {
220220
FnOnce,
221221
FnMut,
222222
Fn,
223+
224+
AsyncFnOnce,
225+
AsyncFnMut,
226+
AsyncFn,
223227
}
224228

225229
impl fmt::Display for FnTrait {
@@ -228,6 +232,9 @@ impl fmt::Display for FnTrait {
228232
FnTrait::FnOnce => write!(f, "FnOnce"),
229233
FnTrait::FnMut => write!(f, "FnMut"),
230234
FnTrait::Fn => write!(f, "Fn"),
235+
FnTrait::AsyncFnOnce => write!(f, "AsyncFnOnce"),
236+
FnTrait::AsyncFnMut => write!(f, "AsyncFnMut"),
237+
FnTrait::AsyncFn => write!(f, "AsyncFn"),
231238
}
232239
}
233240
}
@@ -238,6 +245,9 @@ impl FnTrait {
238245
FnTrait::FnOnce => "call_once",
239246
FnTrait::FnMut => "call_mut",
240247
FnTrait::Fn => "call",
248+
FnTrait::AsyncFnOnce => "async_call_once",
249+
FnTrait::AsyncFnMut => "async_call_mut",
250+
FnTrait::AsyncFn => "async_call",
241251
}
242252
}
243253

@@ -246,6 +256,9 @@ impl FnTrait {
246256
FnTrait::FnOnce => LangItem::FnOnce,
247257
FnTrait::FnMut => LangItem::FnMut,
248258
FnTrait::Fn => LangItem::Fn,
259+
FnTrait::AsyncFnOnce => LangItem::AsyncFnOnce,
260+
FnTrait::AsyncFnMut => LangItem::AsyncFnMut,
261+
FnTrait::AsyncFn => LangItem::AsyncFn,
249262
}
250263
}
251264

@@ -254,15 +267,19 @@ impl FnTrait {
254267
LangItem::FnOnce => Some(FnTrait::FnOnce),
255268
LangItem::FnMut => Some(FnTrait::FnMut),
256269
LangItem::Fn => Some(FnTrait::Fn),
270+
LangItem::AsyncFnOnce => Some(FnTrait::AsyncFnOnce),
271+
LangItem::AsyncFnMut => Some(FnTrait::AsyncFnMut),
272+
LangItem::AsyncFn => Some(FnTrait::AsyncFn),
257273
_ => None,
258274
}
259275
}
260276

261277
pub const fn to_chalk_ir(self) -> rust_ir::ClosureKind {
278+
// Chalk doesn't support async fn traits.
262279
match self {
263-
FnTrait::FnOnce => rust_ir::ClosureKind::FnOnce,
264-
FnTrait::FnMut => rust_ir::ClosureKind::FnMut,
265-
FnTrait::Fn => rust_ir::ClosureKind::Fn,
280+
FnTrait::AsyncFnOnce | FnTrait::FnOnce => rust_ir::ClosureKind::FnOnce,
281+
FnTrait::AsyncFnMut | FnTrait::FnMut => rust_ir::ClosureKind::FnMut,
282+
FnTrait::AsyncFn | FnTrait::Fn => rust_ir::ClosureKind::Fn,
266283
}
267284
}
268285

@@ -271,6 +288,9 @@ impl FnTrait {
271288
FnTrait::FnOnce => Name::new_symbol_root(sym::call_once.clone()),
272289
FnTrait::FnMut => Name::new_symbol_root(sym::call_mut.clone()),
273290
FnTrait::Fn => Name::new_symbol_root(sym::call.clone()),
291+
FnTrait::AsyncFnOnce => Name::new_symbol_root(sym::async_call_once.clone()),
292+
FnTrait::AsyncFnMut => Name::new_symbol_root(sym::async_call_mut.clone()),
293+
FnTrait::AsyncFn => Name::new_symbol_root(sym::async_call.clone()),
274294
}
275295
}
276296

‎crates/ide-diagnostics/src/handlers/expected_function.rs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,4 +37,25 @@ fn foo() {
3737
"#,
3838
);
3939
}
40+
41+
#[test]
42+
fn no_error_for_async_fn_traits() {
43+
check_diagnostics(
44+
r#"
45+
//- minicore: async_fn
46+
async fn f(it: impl AsyncFn(u32) -> i32) {
47+
let fut = it(0);
48+
let _: i32 = fut.await;
49+
}
50+
async fn g(mut it: impl AsyncFnMut(u32) -> i32) {
51+
let fut = it(0);
52+
let _: i32 = fut.await;
53+
}
54+
async fn h(it: impl AsyncFnOnce(u32) -> i32) {
55+
let fut = it(0);
56+
let _: i32 = fut.await;
57+
}
58+
"#,
59+
);
60+
}
4061
}

‎crates/intern/src/symbol/symbols.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,9 @@ define_symbols! {
150150
C,
151151
call_mut,
152152
call_once,
153+
async_call_once,
154+
async_call_mut,
155+
async_call,
153156
call,
154157
cdecl,
155158
Center,
@@ -221,6 +224,9 @@ define_symbols! {
221224
fn_mut,
222225
fn_once_output,
223226
fn_once,
227+
async_fn_once,
228+
async_fn_mut,
229+
async_fn,
224230
fn_ptr_addr,
225231
fn_ptr_trait,
226232
format_alignment,
@@ -334,6 +340,8 @@ define_symbols! {
334340
Option,
335341
Ord,
336342
Output,
343+
CallRefFuture,
344+
CallOnceFuture,
337345
owned_box,
338346
packed,
339347
panic_2015,

0 commit comments

Comments
 (0)