Add FnOnce trait, and provide impl for Function type

book/src/clauses/well_known_traits.md

@@ -28,7 +28,7 @@ Some common examples of auto traits are `Send` and `Sync`.
 [coinductive_section]: ../engine/logic/coinduction.html#coinduction-and-refinement-strands
 
 # Current state 
-| Type            | Copy | Clone | Sized | Unsize | Drop | Fn  | Unpin  | Generator | auto traits |
+| Type            | Copy | Clone | Sized | Unsize | Drop | FnOnce/FnMut/Fn  | Unpin  | Generator | auto traits |
 | ---             | ---  | ---   | ---   | ---    | ---  | --- | ---    |  ---      |  ---        |
 | tuple types     | ✅    | ✅    | ✅     | ✅     | ⚬    | ⚬  |  ⚬      |  ⚬       |   ❌         |
 | structs         | ⚬    | ⚬    |  ✅    | ✅     | ⚬    | ⚬  |  ⚬      |  ⚬       |   ✅         |
@@ -37,7 +37,7 @@ Some common examples of auto traits are `Send` and `Sync`.
 | never type      | 📚   |  📚   |  ✅   |  ⚬    |  ⚬    | ⚬   |   ⚬    |  ⚬       |   ❌       |
 | trait objects   | ⚬    | ⚬    | ⚬     |  ✅    | ⚬    | ⚬   | ⚬      |  ⚬       |    ⚬        |
 | functions defs  | ✅    | ✅    | ✅     | ⚬     | ⚬    |  ❌  | ⚬      |  ⚬       |    ❌         |
-| functions ptrs  | ✅    | ✅    | ✅     | ⚬     | ⚬    | ❌   | ⚬      |  ⚬       |    ❌         |
+| functions ptrs  | ✅    | ✅    | ✅     | ⚬     | ⚬    |  ✅  | ⚬      |  ⚬       |    ❌         |
 | raw ptrs        | 📚   |  📚  |   ✅   |  ⚬    |   ⚬   |  ⚬  |   ⚬    |   ⚬      |      ❌      |
 | immutable refs  | 📚   |  📚  |   ✅   |  ⚬    |   ⚬   |  ⚬  |   ⚬    |   ⚬      |      ❌      |
 | mutable refs    | ⚬    |  ⚬   |   ✅   |  ⚬    |   ⚬   |  ⚬  |   ⚬    |   ⚬      |      ❌      |
@@ -55,4 +55,4 @@ legend:
 📚 - implementation provided in libcore  
 ❌ - not implemented  
 
-❌ after a type name means that type is not yet in chalk
+❌ after a type name means that type is not yet in chalk

chalk-integration/src/lowering.rs

@@ -1394,18 +1394,20 @@ impl LowerTy for Ty {
 
             Ty::ForAll {
                 ref lifetime_names,
-                ref ty,
+                ref types,
             } => {
                 let quantified_env = env.introduce(lifetime_names.iter().map(|id| {
                     chalk_ir::WithKind::new(chalk_ir::VariableKind::Lifetime, id.str.clone())
                 }))?;
 
+                let mut lowered_tys = Vec::with_capacity(types.len());
+                for ty in types {
+                    lowered_tys.push(ty.lower(&quantified_env)?.cast(interner));
+                }
+
                 let function = chalk_ir::Fn {
                     num_binders: lifetime_names.len(),
-                    substitution: Substitution::from(
-                        interner,
-                        Some(ty.lower(&quantified_env)?.cast(interner)),
-                    ),
+                    substitution: Substitution::from(interner, lowered_tys),
                 };
                 Ok(chalk_ir::TyData::Function(function).intern(interner))
             }
@@ -1822,6 +1824,9 @@ impl LowerWellKnownTrait for WellKnownTrait {
             Self::Copy => rust_ir::WellKnownTrait::Copy,
             Self::Clone => rust_ir::WellKnownTrait::Clone,
             Self::Drop => rust_ir::WellKnownTrait::Drop,
+            Self::FnOnce => rust_ir::WellKnownTrait::FnOnce,
+            Self::FnMut => rust_ir::WellKnownTrait::FnMut,
+            Self::Fn => rust_ir::WellKnownTrait::Fn,
         }
     }
 }

chalk-parse/src/ast.rs

@@ -68,6 +68,9 @@ pub enum WellKnownTrait {
     Copy,
     Clone,
     Drop,
+    FnOnce,
+    FnMut,
+    Fn,
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -210,7 +213,7 @@ pub enum Ty {
     },
     ForAll {
         lifetime_names: Vec<Identifier>,
-        ty: Box<Ty>,
+        types: Vec<Box<Ty>>,
     },
     Tuple {
         types: Vec<Box<Ty>>,

chalk-parse/src/parser.lalrpop

@@ -51,6 +51,9 @@ WellKnownTrait: WellKnownTrait = {
      "#" "[" "lang" "(" "copy" ")" "]" => WellKnownTrait::Copy,
      "#" "[" "lang" "(" "clone" ")" "]" => WellKnownTrait::Clone,
      "#" "[" "lang" "(" "drop" ")" "]" => WellKnownTrait::Drop,
+     "#" "[" "lang" "(" "fn_once" ")" "]" => WellKnownTrait::FnOnce,
+     "#" "[" "lang" "(" "fn_mut" ")" "]" => WellKnownTrait::FnMut,
+     "#" "[" "lang" "(" "fn" ")" "]" => WellKnownTrait::Fn,
 };
 
 StructDefn: StructDefn = {
@@ -220,16 +223,22 @@ pub Ty: Ty = {
 };
 
 TyWithoutId: Ty = {
-    "for" "<" <l:Comma<LifetimeId>> ">" "fn" "(" <t:Ty> ")" => Ty::ForAll {
+    "for" "<" <l:Comma<LifetimeId>> ">" "fn" "(" <types:Comma<Ty>> ")" <ret_ty:FnReturn?> => Ty::ForAll {
         lifetime_names: l,
-        ty: Box::new(t)
+        types: types
+                   .into_iter()
+                   .chain(std::iter::once(ret_ty.unwrap_or_else(|| Ty::Tuple { types: Vec::new() })))
+                   .map(Box::new).collect()
     },
     <ScalarType> => Ty::Scalar { ty: <> },
     "str" => Ty::Str,
     "!" => Ty::Never,
-    "fn" "(" <t:Ty> ")" => Ty::ForAll {
+    "fn" "(" <types:Comma<Ty>> ")" <ret_ty:FnReturn?> => Ty::ForAll {
         lifetime_names: vec![],
-        ty: Box::new(t)
+        types: types
+                   .into_iter()
+                   .chain(std::iter::once(ret_ty.unwrap_or_else(|| Ty::Tuple { types: Vec::new() })))
+                   .map(Box::new).collect()
     },
     "dyn" <b:Plus<QuantifiedInlineBound>> "+" <l:Lifetime> => Ty::Dyn {
         bounds: b,

chalk-solve/src/clauses.rs

@@ -358,15 +358,21 @@ fn program_clauses_that_could_match<I: Interner>(
 
                 let trait_datum = db.trait_datum(trait_id);
 
+                let self_ty = alias.self_type_parameter(interner);
+
                 // Flounder if the self-type is unknown and the trait is non-enumerable.
                 //
                 // e.g., Normalize(<?X as Iterator>::Item = u32)
-                if (alias.self_type_parameter(interner).is_var(interner))
-                    && trait_datum.is_non_enumerable_trait()
-                {
+                if (self_ty.is_var(interner)) && trait_datum.is_non_enumerable_trait() {
                     return Err(Floundered);
                 }
 
+                if let Some(well_known) = trait_datum.well_known {
+                    builtin_traits::add_builtin_assoc_program_clauses(
+                        db, builder, well_known, self_ty,
+                    );
+                }
+
                 push_program_clauses_for_associated_type_values_in_impls_of(
                     builder,
                     trait_id,

chalk-solve/src/clauses/builtin_traits.rs

@@ -4,6 +4,7 @@ use chalk_ir::{Substitution, Ty};
 
 mod clone;
 mod copy;
+mod fn_family;
 mod sized;
 
 /// For well known traits we have special hard-coded impls, either as an
@@ -32,12 +33,32 @@ pub fn add_builtin_program_clauses<I: Interner>(
             WellKnownTrait::Sized => sized::add_sized_program_clauses(db, builder, &trait_ref, ty),
             WellKnownTrait::Copy => copy::add_copy_program_clauses(db, builder, &trait_ref, ty),
             WellKnownTrait::Clone => clone::add_clone_program_clauses(db, builder, &trait_ref, ty),
+            WellKnownTrait::FnOnce | WellKnownTrait::FnMut | WellKnownTrait::Fn => {
+                fn_family::add_fn_trait_program_clauses(db, builder, trait_ref.trait_id, self_ty)
+            }
             // Drop impls are provided explicitly
             WellKnownTrait::Drop => (),
         }
     });
 }
 
+/// Like `add_builtin_program_clauses`, but for `DomainGoal::Normalize` involving
+/// a projection (e.g. `<fn(u8) as FnOnce<(u8,)>>::Output`)
+pub fn add_builtin_assoc_program_clauses<I: Interner>(
+    db: &dyn RustIrDatabase<I>,
+    builder: &mut ClauseBuilder<'_, I>,
+    well_known: WellKnownTrait,
+    self_ty: Ty<I>,
+) {
+    match well_known {
+        WellKnownTrait::FnOnce => {
+            let trait_id = db.well_known_trait_id(well_known).unwrap();
+            fn_family::add_fn_trait_program_clauses(db, builder, trait_id, self_ty);
+        }
+        _ => {}
+    }
+}
+
 /// Given a trait ref `T0: Trait` and a list of types `U0..Un`, pushes a clause of the form
 /// `Implemented(T0: Trait) :- Implemented(U0: Trait) .. Implemented(Un: Trait)`
 pub fn needs_impl_for_tys<I: Interner>(

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

@@ -0,0 +1,87 @@
+use crate::clauses::ClauseBuilder;
+use crate::infer::instantiate::IntoBindersAndValue;
+use crate::rust_ir::WellKnownTrait;
+use crate::{Interner, RustIrDatabase, TraitRef};
+use chalk_ir::{
+    AliasTy, ApplicationTy, Binders, Normalize, ProjectionTy, Substitution, TraitId, Ty, TyData,
+    TypeName, VariableKinds,
+};
+
+/// Handles clauses for FnOnce/FnMut/Fn.
+/// If `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
+pub fn add_fn_trait_program_clauses<I: Interner>(
+    db: &dyn RustIrDatabase<I>,
+    builder: &mut ClauseBuilder<'_, I>,
+    trait_id: TraitId<I>,
+    self_ty: Ty<I>,
+) {
+    let interner = db.interner();
+    match self_ty.data(interner) {
+        TyData::Function(fn_val) => {
+            let (binders, orig_sub) = fn_val.into_binders_and_value(interner);
+            let bound_ref = Binders::new(VariableKinds::from(interner, binders), orig_sub);
+            builder.push_binders(&bound_ref, |builder, orig_sub| {
+                // The last parameter represents the function return type
+                let (arg_sub, fn_output_ty) = orig_sub
+                    .parameters(interner)
+                    .split_at(orig_sub.len(interner) - 1);
+                let arg_sub = Substitution::from(interner, arg_sub);
+                let fn_output_ty = fn_output_ty[0].assert_ty_ref(interner);
+
+                // We are constructing a reference to `FnOnce<Args>`, where
+                // `Args` is a tuple of the function's argument types
+                let tupled = Ty::new(
+                    interner,
+                    TyData::Apply(ApplicationTy {
+                        name: TypeName::Tuple(arg_sub.len(interner)),
+                        substitution: arg_sub.clone(),
+                    }),
+                );
+
+                let tupled_sub = Substitution::from(interner, vec![self_ty.clone(), tupled]);
+                // Given a function type `fn(A1, A2, ..., AN)`, construct a `TraitRef`
+                // of the form `fn(A1, A2, ..., AN): FnOnce<(A1, A2, ..., AN)>`
+                let new_trait_ref = TraitRef {
+                    trait_id,
+                    substitution: tupled_sub.clone(),
+                };
+
+                builder.push_fact(new_trait_ref.clone());
+
+                if let Some(WellKnownTrait::FnOnce) = db.trait_datum(trait_id).well_known {
+                    //The `Output` type is defined on the `FnOnce`
+                    let fn_once = db.trait_datum(trait_id);
+                    assert_eq!(fn_once.well_known, Some(WellKnownTrait::FnOnce));
+                    let assoc_types = &fn_once.associated_ty_ids;
+                    assert_eq!(
+                        assoc_types.len(),
+                        1,
+                        "FnOnce trait should have exactly one associated type, found {:?}",
+                        assoc_types
+                    );
+
+                    // Construct `Normalize(<fn(A) -> B as FnOnce<(A,)>>::Output -> B)`
+                    let assoc_output_ty = assoc_types[0];
+                    let proj_ty = ProjectionTy {
+                        associated_ty_id: assoc_output_ty,
+                        substitution: tupled_sub,
+                    };
+                    let normalize = Normalize {
+                        alias: AliasTy::Projection(proj_ty),
+                        ty: fn_output_ty.clone(),
+                    };
+
+                    builder.push_fact(normalize);
+                }
+            });
+        }
+        _ => {}
+    }
+}

chalk-solve/src/rust_ir.rs

@@ -216,6 +216,12 @@ pub enum WellKnownTrait {
     Copy,
     Clone,
     Drop,
+    /// The trait `FnOnce<Args>` - the generic argument `Args` is always a tuple
+    /// corresponding to the arguments of a function implementing this trait.
+    /// E.g. `fn(u8, bool): FnOnce<(u8, bool)>`
+    FnOnce,
+    FnMut,
+    Fn,
 }
 
 impl<I: Interner> TraitDatum<I> {

chalk-solve/src/wf.rs

@@ -471,7 +471,12 @@ impl WfWellKnownGoals {
     ) -> Option<Goal<I>> {
         match db.trait_datum(trait_ref.trait_id).well_known? {
             WellKnownTrait::Copy => Self::copy_impl_constraint(db, trait_ref),
-            WellKnownTrait::Drop | WellKnownTrait::Clone | WellKnownTrait::Sized => None,
+            WellKnownTrait::Drop
+            | WellKnownTrait::Clone
+            | WellKnownTrait::Sized
+            | WellKnownTrait::FnOnce
+            | WellKnownTrait::FnMut
+            | WellKnownTrait::Fn => None,
         }
     }
 
@@ -486,7 +491,10 @@ impl WfWellKnownGoals {
 
         match db.trait_datum(impl_datum.trait_id()).well_known? {
             // You can't add a manual implementation of Sized
-            WellKnownTrait::Sized => Some(GoalData::CannotProve(()).intern(interner)),
+            WellKnownTrait::Sized
+            | WellKnownTrait::FnOnce
+            | WellKnownTrait::FnMut
+            | WellKnownTrait::Fn => Some(GoalData::CannotProve(()).intern(interner)),
             WellKnownTrait::Drop => Self::drop_impl_constraint(db, impl_datum),
             WellKnownTrait::Copy | WellKnownTrait::Clone => None,
         }