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

@@ -295,7 +295,7 @@ fn program_clauses_that_could_match<I: Interner>(
             }
 
             if let Some(well_known) = trait_datum.well_known {
-                builtin_traits::add_builtin_program_clauses(db, builder, well_known, trait_ref);
+                builtin_traits::add_builtin_program_clauses(db, builder, well_known, trait_ref)?;
             }
         }
         DomainGoal::Holds(WhereClause::AliasEq(alias_eq)) => match &alias_eq.alias {
@@ -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/builder.rs

@@ -114,7 +114,11 @@ impl<'me, I: Interner> ClauseBuilder<'me, I> {
     /// The new binders are always pushed onto the end of the internal
     /// list of binders; this means that any extant values where were
     /// created referencing the *old* list of binders are still valid.
-    pub fn push_binders<V>(&mut self, binders: &Binders<V>, op: impl FnOnce(&mut Self, V::Result))
+    pub fn push_binders<R, V>(
+        &mut self,
+        binders: &Binders<V>,
+        op: impl FnOnce(&mut Self, V::Result) -> R,
+    ) -> R
     where
         V: Fold<I> + HasInterner<Interner = I>,
         V::Result: std::fmt::Debug,
@@ -134,10 +138,11 @@ impl<'me, I: Interner> ClauseBuilder<'me, I> {
 
         let value = binders.substitute(self.interner(), &self.parameters[old_len..]);
         debug!("push_binders: value={:?}", value);
-        op(self, value);
+        let res = op(self, value);
 
         self.binders.truncate(old_len);
         self.parameters.truncate(old_len);
+        res
     }
 
     /// Push a single binder, for a type, at the end of the binder

chalk-solve/src/clauses/builtin_traits.rs

@@ -1,9 +1,10 @@
 use super::{builder::ClauseBuilder, generalize};
 use crate::{Interner, RustIrDatabase, TraitRef, WellKnownTrait};
-use chalk_ir::{Substitution, Ty};
+use chalk_ir::{Floundered, Substitution, Ty};
 
 mod clone;
 mod copy;
+mod fn_family;
 mod sized;
 
 /// For well known traits we have special hard-coded impls, either as an
@@ -13,7 +14,7 @@ pub fn add_builtin_program_clauses<I: Interner>(
     builder: &mut ClauseBuilder<'_, I>,
     well_known: WellKnownTrait,
     trait_ref: &TraitRef<I>,
-) {
+) -> Result<(), Floundered> {
     // If `trait_ref` contains bound vars, we want to universally quantify them.
     // `Generalize` collects them for us.
     let generalized = generalize::Generalize::apply(db.interner(), trait_ref);
@@ -25,17 +26,39 @@ pub fn add_builtin_program_clauses<I: Interner>(
             if force_impl {
                 builder.push_fact(trait_ref.clone());
             }
-            return;
+            return Ok(());
         }
 
         match well_known {
             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 => (),
         }
-    });
+        Ok(())
+    })
+}
+
+/// 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>,
+) -> Result<(), Floundered> {
+    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)?;
+        }
+        _ => {}
+    }
+    Ok(())
 }
 
 /// Given a trait ref `T0: Trait` and a list of types `U0..Un`, pushes a clause of the form

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

@@ -0,0 +1,90 @@
+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, Floundered, 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>,
+) -> Result<(), Floundered> {
+    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);
+                }
+            });
+            Ok(())
+        }
+        // Function traits are non-enumerable
+        TyData::InferenceVar(..) | TyData::Alias(..) => Err(Floundered),
+        _ => Ok(()),
+    }
+}

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,
         }