Skip to content

Commit 7baa85d

Browse files
ChayimFriedman2ChayimFriedman2
committed
Fix inference of AsyncFnX return type
1 parent 2e9c2ce commit 7baa85d

File tree

5 files changed

+81
-11
lines changed

5 files changed

+81
-11
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,7 @@ impl chalk_solve::RustIrDatabase<Interner> for ChalkContext<'_> {
259259
}
260260
fn well_known_trait_id(
261261
&self,
262-
well_known_trait: rust_ir::WellKnownTrait,
262+
well_known_trait: WellKnownTrait,
263263
) -> Option<chalk_ir::TraitId<Interner>> {
264264
let lang_attr = lang_item_from_well_known_trait(well_known_trait);
265265
let trait_ = lang_attr.resolve_trait(self.db, self.krate)?;

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1463,6 +1463,8 @@ impl HirDisplay for Ty {
14631463
}
14641464
if f.closure_style == ClosureStyle::RANotation || !sig.ret().is_unit() {
14651465
write!(f, " -> ")?;
1466+
// FIXME: We display `AsyncFn` as `-> impl Future`, but this is hard to fix because
1467+
// we don't have a trait environment here, required to normalize `<Ret as Future>::Output`.
14661468
sig.ret().hir_fmt(f)?;
14671469
}
14681470
} else {

‎src/tools/rust-analyzer/crates/hir-ty/src/infer/closure.rs

Lines changed: 46 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ use crate::{
3838
infer::{BreakableKind, CoerceMany, Diverges, coerce::CoerceNever},
3939
make_binders,
4040
mir::{BorrowKind, MirSpan, MutBorrowKind, ProjectionElem},
41-
to_chalk_trait_id,
41+
to_assoc_type_id, to_chalk_trait_id,
4242
traits::FnTrait,
4343
utils::{self, elaborate_clause_supertraits},
4444
};
@@ -245,7 +245,7 @@ impl InferenceContext<'_> {
245245
}
246246

247247
fn deduce_closure_kind_from_predicate_clauses(
248-
&self,
248+
&mut self,
249249
expected_ty: &Ty,
250250
clauses: impl DoubleEndedIterator<Item = WhereClause>,
251251
closure_kind: ClosureKind,
@@ -378,7 +378,7 @@ impl InferenceContext<'_> {
378378
}
379379

380380
fn deduce_sig_from_projection(
381-
&self,
381+
&mut self,
382382
closure_kind: ClosureKind,
383383
projection_ty: &ProjectionTy,
384384
projected_ty: &Ty,
@@ -392,13 +392,16 @@ impl InferenceContext<'_> {
392392

393393
// For now, we only do signature deduction based off of the `Fn` and `AsyncFn` traits,
394394
// for closures and async closures, respectively.
395-
match closure_kind {
396-
ClosureKind::Closure | ClosureKind::Async
397-
if self.fn_trait_kind_from_trait_id(trait_).is_some() =>
398-
{
399-
self.extract_sig_from_projection(projection_ty, projected_ty)
400-
}
401-
_ => None,
395+
let fn_trait_kind = self.fn_trait_kind_from_trait_id(trait_)?;
396+
if !matches!(closure_kind, ClosureKind::Closure | ClosureKind::Async) {
397+
return None;
398+
}
399+
if fn_trait_kind.is_async() {
400+
// If the expected trait is `AsyncFn(...) -> X`, we don't know what the return type is,
401+
// but we do know it must implement `Future<Output = X>`.
402+
self.extract_async_fn_sig_from_projection(projection_ty, projected_ty)
403+
} else {
404+
self.extract_sig_from_projection(projection_ty, projected_ty)
402405
}
403406
}
404407

@@ -424,6 +427,39 @@ impl InferenceContext<'_> {
424427
)))
425428
}
426429

430+
fn extract_async_fn_sig_from_projection(
431+
&mut self,
432+
projection_ty: &ProjectionTy,
433+
projected_ty: &Ty,
434+
) -> Option<FnSubst<Interner>> {
435+
let arg_param_ty = projection_ty.substitution.as_slice(Interner)[1].assert_ty_ref(Interner);
436+
437+
let TyKind::Tuple(_, input_tys) = arg_param_ty.kind(Interner) else {
438+
return None;
439+
};
440+
441+
let ret_param_future_output = projected_ty;
442+
let ret_param_future = self.table.new_type_var();
443+
let future_output =
444+
LangItem::FutureOutput.resolve_type_alias(self.db, self.resolver.krate())?;
445+
let future_projection = crate::AliasTy::Projection(crate::ProjectionTy {
446+
associated_ty_id: to_assoc_type_id(future_output),
447+
substitution: Substitution::from1(Interner, ret_param_future.clone()),
448+
});
449+
self.table.register_obligation(
450+
crate::AliasEq { alias: future_projection, ty: ret_param_future_output.clone() }
451+
.cast(Interner),
452+
);
453+
454+
Some(FnSubst(Substitution::from_iter(
455+
Interner,
456+
input_tys.iter(Interner).map(|t| t.cast(Interner)).chain(Some(GenericArg::new(
457+
Interner,
458+
chalk_ir::GenericArgData::Ty(ret_param_future),
459+
))),
460+
)))
461+
}
462+
427463
fn fn_trait_kind_from_trait_id(&self, trait_id: hir_def::TraitId) -> Option<FnTrait> {
428464
FnTrait::from_lang_item(self.db.lang_attr(trait_id.into())?)
429465
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4903,3 +4903,30 @@ fn main() {
49034903
"#]],
49044904
);
49054905
}
4906+
4907+
#[test]
4908+
fn async_fn_return_type() {
4909+
check_infer(
4910+
r#"
4911+
//- minicore: async_fn
4912+
fn foo<F: AsyncFn() -> R, R>(_: F) -> R {
4913+
loop {}
4914+
}
4915+
4916+
fn main() {
4917+
foo(async move || ());
4918+
}
4919+
"#,
4920+
expect![[r#"
4921+
29..30 '_': F
4922+
40..55 '{ loop {} }': R
4923+
46..53 'loop {}': !
4924+
51..53 '{}': ()
4925+
67..97 '{ ...()); }': ()
4926+
73..76 'foo': fn foo<impl AsyncFn() -> impl Future<Output = ()>, ()>(impl AsyncFn() -> impl Future<Output = ()>)
4927+
73..94 'foo(as...|| ())': ()
4928+
77..93 'async ... || ()': impl AsyncFn() -> impl Future<Output = ()>
4929+
91..93 '()': ()
4930+
"#]],
4931+
);
4932+
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -291,4 +291,9 @@ impl FnTrait {
291291
pub fn get_id(self, db: &dyn HirDatabase, krate: Crate) -> Option<TraitId> {
292292
self.lang_item().resolve_trait(db, krate)
293293
}
294+
295+
#[inline]
296+
pub(crate) fn is_async(self) -> bool {
297+
matches!(self, FnTrait::AsyncFn | FnTrait::AsyncFnMut | FnTrait::AsyncFnOnce)
298+
}
294299
}

0 commit comments

Comments
 (0)