Skip to content

Commit 8466efc

Browse files
VeykrilVeykril
authored
Merge pull request rust-lang#18594 from ChayimFriedman2/async-closures
feat: Support `AsyncFnX` traits
2 parents b4dc94d + 879312e commit 8466efc

File tree

9 files changed

+291
-68
lines changed

9 files changed

+291
-68
lines changed

‎src/tools/rust-analyzer/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

‎src/tools/rust-analyzer/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,

‎src/tools/rust-analyzer/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

‎src/tools/rust-analyzer/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) {

‎src/tools/rust-analyzer/crates/hir-ty/src/tests/traits.rs‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4834,3 +4834,53 @@ fn bar(v: *const ()) {
48344834
"#]],
48354835
);
48364836
}
4837+
4838+
#[test]
4839+
fn async_fn_traits() {
4840+
check_infer(
4841+
r#"
4842+
//- minicore: async_fn
4843+
async fn foo<T: AsyncFn(u32) -> i32>(a: T) {
4844+
let fut1 = a(0);
4845+
fut1.await;
4846+
}
4847+
async fn bar<T: AsyncFnMut(u32) -> i32>(mut b: T) {
4848+
let fut2 = b(0);
4849+
fut2.await;
4850+
}
4851+
async fn baz<T: AsyncFnOnce(u32) -> i32>(c: T) {
4852+
let fut3 = c(0);
4853+
fut3.await;
4854+
}
4855+
"#,
4856+
expect![[r#"
4857+
37..38 'a': T
4858+
43..83 '{ ...ait; }': ()
4859+
43..83 '{ ...ait; }': impl Future<Output = ()>
4860+
53..57 'fut1': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4861+
60..61 'a': T
4862+
60..64 'a(0)': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4863+
62..63 '0': u32
4864+
70..74 'fut1': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4865+
70..80 'fut1.await': i32
4866+
124..129 'mut b': T
4867+
134..174 '{ ...ait; }': ()
4868+
134..174 '{ ...ait; }': impl Future<Output = ()>
4869+
144..148 'fut2': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4870+
151..152 'b': T
4871+
151..155 'b(0)': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4872+
153..154 '0': u32
4873+
161..165 'fut2': AsyncFnMut::CallRefFuture<'?, T, (u32,)>
4874+
161..171 'fut2.await': i32
4875+
216..217 'c': T
4876+
222..262 '{ ...ait; }': ()
4877+
222..262 '{ ...ait; }': impl Future<Output = ()>
4878+
232..236 'fut3': AsyncFnOnce::CallOnceFuture<T, (u32,)>
4879+
239..240 'c': T
4880+
239..243 'c(0)': AsyncFnOnce::CallOnceFuture<T, (u32,)>
4881+
241..242 '0': u32
4882+
249..253 'fut3': AsyncFnOnce::CallOnceFuture<T, (u32,)>
4883+
249..259 'fut3.await': i32
4884+
"#]],
4885+
);
4886+
}

‎src/tools/rust-analyzer/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

‎src/tools/rust-analyzer/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
}

‎src/tools/rust-analyzer/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)