Add FnMut and Fn traits

Author: Aaron1011

chalk-integration/src/lowering.rs

@@ -1825,6 +1825,8 @@ impl LowerWellKnownTrait for WellKnownTrait {
             Self::CloneTrait => rust_ir::WellKnownTrait::CloneTrait,
             Self::DropTrait => rust_ir::WellKnownTrait::DropTrait,
             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 {
     CloneTrait,
     DropTrait,
     FnOnceTrait,
+    FnMutTrait,
+    FnTrait
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]

chalk-parse/src/parser.lalrpop

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

chalk-solve/src/clauses/builtin_traits.rs

@@ -39,8 +39,8 @@ pub fn add_builtin_program_clauses<I: Interner>(
             WellKnownTrait::CloneTrait => {
                 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::DropTrait => (),

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

@@ -474,7 +474,9 @@ impl WfWellKnownGoals {
             WellKnownTrait::DropTrait
             | WellKnownTrait::CloneTrait
             | WellKnownTrait::SizedTrait
-            | WellKnownTrait::FnOnceTrait => None,
+            | WellKnownTrait::FnOnceTrait
+            | WellKnownTrait::FnMutTrait
+            | WellKnownTrait::FnTrait => None,
         }
     }
 
@@ -488,9 +490,12 @@ impl WfWellKnownGoals {
         let interner = db.interner();
 
         match db.trait_datum(impl_datum.trait_id()).well_known? {
-            // You can't add a manual implementation of Sized or FnOnce
-            WellKnownTrait::SizedTrait | WellKnownTrait::FnOnceTrait => {
-                Some(GoalData::CannotProve(()).intern(interner))
+            // You can't add a manual implementation of Sized or FnOnce/FnMut/Fn
+            WellKnownTrait::SizedTrait
+                | WellKnownTrait::FnOnceTrait
+                | WellKnownTrait::FnMutTrait
+                | WellKnownTrait::FnTrait => {
+                    Some(GoalData::CannotProve(()).intern(interner))
             }
             WellKnownTrait::DropTrait => Self::drop_impl_constraint(db, impl_datum),
             WellKnownTrait::CopyTrait | WellKnownTrait::CloneTrait => 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 {