Add FnMut and Fn traits

Author: Aaron1011

chalk-integration/src/lowering.rs

@@ -1825,6 +1825,8 @@ impl LowerWellKnownTrait for WellKnownTrait {
             Self::Clone => rust_ir::WellKnownTrait::Clone,
             Self::Drop => rust_ir::WellKnownTrait::Drop,
             Self::FnOnceTrait => rust_ir::WellKnownTrait::FnOnceTrait,
+            Self::FnMutTrait => rust_ir::WellKnownTrait::FnMutTrait,
+            Self::FnTrait => rust_ir::WellKnownTrait::FnTrait,
         }
     }
 }

chalk-parse/src/ast.rs

@@ -69,6 +69,8 @@ pub enum WellKnownTrait {
     Clone,
     Drop,
     FnOnceTrait,
+    FnMutTrait,
+    FnTrait,
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]

chalk-parse/src/parser.lalrpop

@@ -52,6 +52,8 @@ WellKnownTrait: WellKnownTrait = {
      "#" "[" "lang" "(" "clone" ")" "]" => WellKnownTrait::Clone,
      "#" "[" "lang" "(" "drop" ")" "]" => WellKnownTrait::Drop,
      "#" "[" "lang" "(" "fn_once" ")" "]" => WellKnownTrait::FnOnceTrait,
+     "#" "[" "lang" "(" "fn_mut" ")" "]" => WellKnownTrait::FnMutTrait,
+     "#" "[" "lang" "(" "fn" ")" "]" => WellKnownTrait::FnTrait,
 };
 
 StructDefn: StructDefn = {

chalk-solve/src/clauses/builtin_traits.rs

@@ -33,8 +33,8 @@ 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::FnOnceTrait => {
-                fn_::add_fn_once_program_clauses(db, builder, &trait_ref, ty)
+            WellKnownTrait::FnOnceTrait | WellKnownTrait::FnMutTrait | WellKnownTrait::FnTrait => {
+                fn_::add_fn_trait_program_clauses(db, builder, &trait_ref, ty)
             }
             // Drop impls are provided explicitly
             WellKnownTrait::Drop => (),

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

@@ -3,7 +3,8 @@ use crate::infer::instantiate::IntoBindersAndValue;
 use crate::{Interner, RustIrDatabase, TraitRef};
 use chalk_ir::{ApplicationTy, Binders, Substitution, Ty, TyData, TypeName, VariableKinds};
 
-pub fn add_fn_once_program_clauses<I: Interner>(
+// Handles clauses for FnOnce/FnMut/Fn
+pub fn add_fn_trait_program_clauses<I: Interner>(
     db: &dyn RustIrDatabase<I>,
     builder: &mut ClauseBuilder<'_, I>,
     trait_ref: &TraitRef<I>,

chalk-solve/src/rust_ir.rs

@@ -220,6 +220,8 @@ pub enum WellKnownTrait {
     /// corresponding to the arguments of a function implementing this trait.
     /// E.g. `fn(u8, bool): FnOnce<(u8, bool)>`
     FnOnceTrait,
+    FnMutTrait,
+    FnTrait,
 }
 
 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 | WellKnownTrait::FnOnceTrait => None,
+            WellKnownTrait::Drop
+            | WellKnownTrait::Clone
+            | WellKnownTrait::Sized
+            | WellKnownTrait::FnOnceTrait
+            | WellKnownTrait::FnMutTrait
+            | WellKnownTrait::FnTrait => 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 | WellKnownTrait::FnOnceTrait => Some(GoalData::CannotProve(()).intern(interner)),
+            WellKnownTrait::Sized
+            | WellKnownTrait::FnOnceTrait
+            | WellKnownTrait::FnMutTrait
+            | WellKnownTrait::FnTrait => Some(GoalData::CannotProve(()).intern(interner)),
             WellKnownTrait::Drop => Self::drop_impl_constraint(db, impl_datum),
             WellKnownTrait::Copy | WellKnownTrait::Clone => None,
         }

tests/test/functions.rs

@@ -6,6 +6,12 @@ fn function_implement_fn_once() {
         program {
             #[lang(fn_once)]
             trait FnOnce<Args> { }
+
+            #[lang(fn_mut)]
+            trait FnMut<Args> where Self: FnOnce<Args> { }
+
+            #[lang(fn)]
+            trait Fn<Args> where Self: FnMut<Args> { }
         }
 
         goal {
@@ -14,6 +20,18 @@ fn function_implement_fn_once() {
             "Unique; substitution [], lifetime constraints []"
         }
 
+        goal {
+            fn(u8): FnMut<(u8,)>
+        } yields {
+            "Unique; substitution [], lifetime constraints []"
+        }
+
+        goal {
+            fn(u8): Fn<(u8,)>
+        } yields {
+            "Unique; substitution [], lifetime constraints []"
+        }
+
         goal {
             fn(u8, u32): FnOnce<(u8,u32)>
         } yields {