Implement built-in supports for AsyncFn* traits

chalk-integration/src/lowering.rs

@@ -1133,6 +1133,9 @@ impl Lower for WellKnownTrait {
             WellKnownTrait::FnOnce => rust_ir::WellKnownTrait::FnOnce,
             WellKnownTrait::FnMut => rust_ir::WellKnownTrait::FnMut,
             WellKnownTrait::Fn => rust_ir::WellKnownTrait::Fn,
+            WellKnownTrait::AsyncFnOnce => rust_ir::WellKnownTrait::AsyncFnOnce,
+            WellKnownTrait::AsyncFnMut => rust_ir::WellKnownTrait::AsyncFnMut,
+            WellKnownTrait::AsyncFn => rust_ir::WellKnownTrait::AsyncFn,
             WellKnownTrait::Unsize => rust_ir::WellKnownTrait::Unsize,
             WellKnownTrait::Unpin => rust_ir::WellKnownTrait::Unpin,
             WellKnownTrait::CoerceUnsized => rust_ir::WellKnownTrait::CoerceUnsized,
@@ -1142,6 +1145,7 @@ impl Lower for WellKnownTrait {
             WellKnownTrait::Tuple => rust_ir::WellKnownTrait::Tuple,
             WellKnownTrait::Pointee => rust_ir::WellKnownTrait::Pointee,
             WellKnownTrait::FnPtr => rust_ir::WellKnownTrait::FnPtr,
+            WellKnownTrait::Future => rust_ir::WellKnownTrait::Future,
         }
     }
 }

chalk-parse/src/ast.rs

@@ -155,6 +155,9 @@ pub enum WellKnownTrait {
     FnOnce,
     FnMut,
     Fn,
+    AsyncFnOnce,
+    AsyncFnMut,
+    AsyncFn,
     Unsize,
     Unpin,
     CoerceUnsized,
@@ -164,6 +167,7 @@ pub enum WellKnownTrait {
     Tuple,
     Pointee,
     FnPtr,
+    Future,
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]

chalk-parse/src/parser.lalrpop

@@ -63,6 +63,9 @@ WellKnownTrait: WellKnownTrait = {
      "#" "[" "lang" "(" "fn_once" ")" "]" => WellKnownTrait::FnOnce,
      "#" "[" "lang" "(" "fn_mut" ")" "]" => WellKnownTrait::FnMut,
      "#" "[" "lang" "(" "fn" ")" "]" => WellKnownTrait::Fn,
+     "#" "[" "lang" "(" "async_fn_once" ")" "]" => WellKnownTrait::AsyncFnOnce,
+     "#" "[" "lang" "(" "async_fn_mut" ")" "]" => WellKnownTrait::AsyncFnMut,
+     "#" "[" "lang" "(" "async_fn" ")" "]" => WellKnownTrait::AsyncFn,
      "#" "[" "lang" "(" "unsize" ")" "]" => WellKnownTrait::Unsize,
      "#" "[" "lang" "(" "unpin" ")" "]" => WellKnownTrait::Unpin,
      "#" "[" "lang" "(" "coerce_unsized" ")" "]" => WellKnownTrait::CoerceUnsized,
@@ -72,6 +75,7 @@ WellKnownTrait: WellKnownTrait = {
      "#" "[" "lang" "(" "tuple_trait" ")" "]" => WellKnownTrait::Tuple,
      "#" "[" "lang" "(" "pointee_trait" ")" "]" => WellKnownTrait::Pointee,
      "#" "[" "lang" "(" "fn_ptr_trait" ")" "]" => WellKnownTrait::FnPtr,
+     "#" "[" "lang" "(" "future" ")" "]" => WellKnownTrait::Future,
 };
 
 AdtReprAttr: AdtReprAttr = {

chalk-solve/src/clauses/builtin_traits.rs

@@ -43,7 +43,12 @@ pub fn add_builtin_program_clauses<I: Interner>(
             WellKnownTrait::Clone => {
                 clone::add_clone_program_clauses(db, builder, trait_ref, ty, binders)?;
             }
-            WellKnownTrait::FnOnce | WellKnownTrait::FnMut | WellKnownTrait::Fn => {
+            WellKnownTrait::FnOnce
+            | WellKnownTrait::FnMut
+            | WellKnownTrait::Fn
+            | WellKnownTrait::AsyncFnOnce
+            | WellKnownTrait::AsyncFnMut
+            | WellKnownTrait::AsyncFn => {
                 fn_family::add_fn_trait_program_clauses(db, builder, well_known, self_ty);
             }
             WellKnownTrait::Unsize => {
@@ -69,7 +74,8 @@ pub fn add_builtin_program_clauses<I: Interner>(
             WellKnownTrait::Unpin
             | WellKnownTrait::Drop
             | WellKnownTrait::CoerceUnsized
-            | WellKnownTrait::DispatchFromDyn => (),
+            | WellKnownTrait::DispatchFromDyn
+            | WellKnownTrait::Future => (),
         }
         Ok(())
     })
@@ -87,7 +93,7 @@ pub fn add_builtin_assoc_program_clauses<I: Interner>(
     // `Generalize` collects them for us.
     let generalized = generalize::Generalize::apply(db.interner(), self_ty);
     builder.push_binders(generalized, |builder, self_ty| match well_known {
-        WellKnownTrait::FnOnce => {
+        WellKnownTrait::FnOnce | WellKnownTrait::AsyncFnOnce => {
             fn_family::add_fn_trait_program_clauses(db, builder, well_known, self_ty);
             Ok(())
         }

chalk-solve/src/clauses/builtin_traits/fn_family.rs

@@ -3,7 +3,7 @@ use crate::rust_ir::{ClosureKind, FnDefInputsAndOutputDatum, WellKnownTrait};
 use crate::{Interner, RustIrDatabase, TraitRef};
 use chalk_ir::cast::Cast;
 use chalk_ir::{
-    AliasTy, Binders, Normalize, ProjectionTy, Safety, Substitution, TraitId, Ty, TyKind,
+    AliasTy, Binders, Goal, Normalize, ProjectionTy, Safety, Substitution, TraitId, Ty, TyKind,
 };
 
 fn push_clauses<I: Interner>(
@@ -19,31 +19,113 @@ fn push_clauses<I: Interner>(
     let tupled = TyKind::Tuple(arg_sub.len(interner), arg_sub).intern(interner);
     let substitution =
         Substitution::from_iter(interner, &[self_ty.cast(interner), tupled.cast(interner)]);
-    builder.push_fact(TraitRef {
-        trait_id,
-        substitution: substitution.clone(),
-    });
 
-    // The `Output` type is defined on the `FnOnce`
-    if let WellKnownTrait::FnOnce = well_known {
-        let trait_datum = db.trait_datum(trait_id);
-        assert_eq!(
-            trait_datum.associated_ty_ids.len(),
-            1,
-            "FnOnce trait should have exactly one associated type, found {:?}",
-            trait_datum.associated_ty_ids
-        );
-        // Constructs the alias. For `Fn`, for example, this would look like
-        // `Normalize(<fn(A) -> B as FnOnce<(A,)>>::Output -> B)`
-        let output_id = trait_datum.associated_ty_ids[0];
-        let alias = AliasTy::Projection(ProjectionTy {
-            associated_ty_id: output_id,
-            substitution,
-        });
-        builder.push_fact(Normalize {
-            alias,
-            ty: return_type,
+    let is_async = matches!(
+        well_known,
+        WellKnownTrait::AsyncFnOnce | WellKnownTrait::AsyncFnMut | WellKnownTrait::AsyncFn
+    );
+
+    if !is_async {
+        builder.push_fact(TraitRef {
+            trait_id,
+            substitution: substitution.clone(),
         });
+
+        // The `Output` type is defined on the `FnOnce`
+        if let WellKnownTrait::FnOnce = well_known {
+            let trait_datum = db.trait_datum(trait_id);
+            assert_eq!(
+                trait_datum.associated_ty_ids.len(),
+                1,
+                "FnOnce trait should have exactly one associated type, found {:?}",
+                trait_datum.associated_ty_ids
+            );
+            // Constructs the alias. For `Fn`, for example, this would look like
+            // `Normalize(<fn(A) -> B as FnOnce<(A,)>>::Output -> B)`
+            let output_id = trait_datum.associated_ty_ids[0];
+            let alias = AliasTy::Projection(ProjectionTy {
+                associated_ty_id: output_id,
+                substitution,
+            });
+            builder.push_fact(Normalize {
+                alias,
+                ty: return_type,
+            });
+        }
+    } else {
+        let sync_counterpart = match well_known {
+            WellKnownTrait::AsyncFnOnce => db.well_known_trait_id(WellKnownTrait::FnOnce).unwrap(),
+            WellKnownTrait::AsyncFnMut => db.well_known_trait_id(WellKnownTrait::FnMut).unwrap(),
+            WellKnownTrait::AsyncFn => db.well_known_trait_id(WellKnownTrait::Fn).unwrap(),
+            _ => unreachable!(),
+        };
+
+        let future = db.well_known_trait_id(WellKnownTrait::Future).unwrap();
+        let sync_counterpart = TraitRef {
+            trait_id: sync_counterpart,
+            substitution: substitution.clone(),
+        };
+        let output_is_future = TraitRef {
+            trait_id: future,
+            substitution: Substitution::from1(interner, return_type.clone()),
+        };
+
+        // This adds the following clause:
+        // `F: AsyncFnX<Arg, Output = O>` :- `F: FnX<Arg, Output: Fut<Output = O>>`
+        // Actually, the `<F as AsyncFnX>::Output = O` part is added in the if let expression below.
+        builder.push_clause(
+            TraitRef {
+                trait_id,
+                substitution: substitution.clone(),
+            },
+            [sync_counterpart.clone(), output_is_future.clone()],
+        );
+
+        if let WellKnownTrait::AsyncFnOnce = well_known {
+            builder.push_bound_ty(|builder, ty| {
+                let trait_datum = db.trait_datum(trait_id);
+                assert_eq!(
+                    trait_datum.associated_ty_ids.len(),
+                    2,
+                    "AsyncFnOnce trait should have exactly two associated types, found {:?}",
+                    trait_datum.associated_ty_ids
+                );
+                let output_id = trait_datum.associated_ty_ids[1];
+                let async_alias = AliasTy::Projection(ProjectionTy {
+                    associated_ty_id: output_id,
+                    substitution,
+                });
+
+                let trait_datum = db.trait_datum(future);
+                assert_eq!(
+                    trait_datum.associated_ty_ids.len(),
+                    1,
+                    "Future trait should have exactly one associated type, found {:?}",
+                    trait_datum.associated_ty_ids
+                );
+                let output_id = trait_datum.associated_ty_ids[0];
+                let future_alias = AliasTy::Projection(ProjectionTy {
+                    associated_ty_id: output_id,
+                    substitution: Substitution::from1(interner, return_type),
+                });
+
+                builder.push_clause(
+                    Normalize {
+                        alias: async_alias,
+                        ty: ty.clone(),
+                    },
+                    [
+                        sync_counterpart.cast::<Goal<_>>(interner),
+                        output_is_future.cast(interner),
+                        Normalize {
+                            alias: future_alias,
+                            ty,
+                        }
+                        .cast(interner),
+                    ],
+                );
+            });
+        }
     }
 }
 
@@ -71,15 +153,20 @@ fn push_clauses_for_apply<I: Interner>(
     });
 }
 
-/// Handles clauses for FnOnce/FnMut/Fn.
-/// If `self_ty` is a function, we push a clause of the form
+/// Handles clauses for FnOnce/FnMut/Fn and AsyncFnOnce/AsyncFnMut/AsyncFn.
+/// For sync traits, `self_ty` is a function, we push a clause of the form
 /// `fn(A1, A2, ..., AN) -> O: FnTrait<(A1, A2, ..., AN)>`, where `FnTrait`
 /// is the trait corresponding to `trait_id` (FnOnce/FnMut/Fn)
 ///
 /// If `trait_id` is `FnOnce`, we also push a clause for the output type of the form:
 /// `Normalize(<fn(A) -> B as FnOnce<(A,)>>::Output -> B)`
 /// We do not add the usual `Implemented(fn(A) -> b as FnOnce<(A,)>` clause
 /// as a condition, since we already called `push_fact` with it
+///
+/// For async traits, we push a clause of the form
+/// `F: AsyncFnX<Arg, Output = O>` :- `F: FnX<Arg, Output: Fut<Output = O>>`,
+/// which corresponds to the implementation
+/// `impl<F, Arg, Fut, O> AsyncFn<A> for F where F: Fn<Arg, Output = Fut>, Fut: Future<Output = O>`.
 pub fn add_fn_trait_program_clauses<I: Interner>(
     db: &dyn RustIrDatabase<I>,
     builder: &mut ClauseBuilder<'_, I>,
@@ -111,9 +198,13 @@ pub fn add_fn_trait_program_clauses<I: Interner>(
             let closure_kind = db.closure_kind(*closure_id, substitution);
             let trait_matches = matches!(
                 (well_known, closure_kind),
-                (WellKnownTrait::Fn, ClosureKind::Fn)
-                    | (WellKnownTrait::FnMut, ClosureKind::FnMut | ClosureKind::Fn)
-                    | (WellKnownTrait::FnOnce, _)
+                (
+                    WellKnownTrait::Fn | WellKnownTrait::AsyncFn,
+                    ClosureKind::Fn
+                ) | (
+                    WellKnownTrait::FnMut | WellKnownTrait::AsyncFnMut,
+                    ClosureKind::FnMut | ClosureKind::Fn
+                ) | (WellKnownTrait::FnOnce | WellKnownTrait::AsyncFnOnce, _)
             );
             if !trait_matches {
                 return;

chalk-solve/src/display/items.rs

@@ -199,6 +199,9 @@ impl<I: Interner> RenderAsRust<I> for TraitDatum<I> {
                 WellKnownTrait::FnOnce => "fn_once",
                 WellKnownTrait::FnMut => "fn_mut",
                 WellKnownTrait::Fn => "fn",
+                WellKnownTrait::AsyncFnOnce => "async_fn_once",
+                WellKnownTrait::AsyncFnMut => "async_fn_mut",
+                WellKnownTrait::AsyncFn => "async_fn",
                 WellKnownTrait::Unsize => "unsize",
                 WellKnownTrait::Unpin => "unpin",
                 WellKnownTrait::CoerceUnsized => "coerce_unsized",
@@ -208,6 +211,7 @@ impl<I: Interner> RenderAsRust<I> for TraitDatum<I> {
                 WellKnownTrait::Tuple => "tuple_trait",
                 WellKnownTrait::Pointee => "pointee",
                 WellKnownTrait::FnPtr => "fn_ptr_trait",
+                WellKnownTrait::Future => "future",
             };
             writeln!(f, "#[lang({})]", name)?;
         }

chalk-solve/src/rust_ir.rs

@@ -270,6 +270,9 @@ pub enum WellKnownTrait {
     FnOnce,
     FnMut,
     Fn,
+    AsyncFnOnce,
+    AsyncFnMut,
+    AsyncFn,
     Unsize,
     Unpin,
     CoerceUnsized,
@@ -279,6 +282,7 @@ pub enum WellKnownTrait {
     Tuple,
     Pointee,
     FnPtr,
+    Future,
 }
 
 chalk_ir::const_visit!(WellKnownTrait);

chalk-solve/src/wf.rs

@@ -427,11 +427,14 @@ where
                     &impl_datum,
                 )
             }
-            WellKnownTrait::Clone | WellKnownTrait::Unpin => true,
+            WellKnownTrait::Clone | WellKnownTrait::Unpin | WellKnownTrait::Future => true,
             // You can't add a manual implementation for the following traits:
             WellKnownTrait::Fn
             | WellKnownTrait::FnOnce
             | WellKnownTrait::FnMut
+            | WellKnownTrait::AsyncFn
+            | WellKnownTrait::AsyncFnOnce
+            | WellKnownTrait::AsyncFnMut
             | WellKnownTrait::Unsize
             | WellKnownTrait::Sized
             | WellKnownTrait::DiscriminantKind