Handle FnOnce::Output

Author: Aaron1011

chalk-parse/src/parser.lalrpop

@@ -223,16 +223,22 @@ pub Ty: Ty = {
 };
 
 TyWithoutId: Ty = {
-    "for" "<" <l:Comma<LifetimeId>> ">" "fn" "(" <types:Comma<Ty>> ")" => Ty::ForAll {
+    "for" "<" <l:Comma<LifetimeId>> ">" "fn" "(" <types:Comma<Ty>> ")" <ret_ty:FnReturn?> => Ty::ForAll {
         lifetime_names: l,
-        types: types.into_iter().map(Box::new).collect()
+        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" "(" <types:Comma<Ty>> ")" => Ty::ForAll {
+    "fn" "(" <types:Comma<Ty>> ")" <ret_ty:FnReturn?> => Ty::ForAll {
         lifetime_names: vec![],
-        types: types.into_iter().map(Box::new).collect()
+        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

@@ -367,6 +367,12 @@ fn program_clauses_that_could_match<I: Interner>(
                     return Err(Floundered);
                 }
 
+                if let Some(well_known) = trait_datum.well_known {
+                    builtin_traits::add_builtin_assoc_program_clauses(
+                        db, builder, well_known, proj,
+                    );
+                }
+
                 push_program_clauses_for_associated_type_values_in_impls_of(
                     builder,
                     trait_id,

chalk-solve/src/clauses/builtin_traits.rs

@@ -1,6 +1,6 @@
 use super::{builder::ClauseBuilder, generalize};
 use crate::{Interner, RustIrDatabase, TraitRef, WellKnownTrait};
-use chalk_ir::{Substitution, Ty};
+use chalk_ir::{ProjectionTy, Substitution, Ty};
 
 mod clone;
 mod copy;
@@ -40,14 +40,37 @@ pub fn add_builtin_program_clauses<I: Interner>(
                 clone::add_clone_program_clauses(db, builder, &trait_ref, ty)
             }
             WellKnownTrait::FnOnceTrait | WellKnownTrait::FnMutTrait | WellKnownTrait::FnTrait => {
-                fn_::add_fn_trait_program_clauses(db, builder, &trait_ref, ty)
+                fn_::add_fn_trait_program_clauses(db, builder, trait_ref.trait_id, ty, false)
             }
             // Drop impls are provided explicitly
             WellKnownTrait::DropTrait => (),
         }
     });
 }
 
+/// 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,
+    proj: &ProjectionTy<I>,
+) {
+    match well_known {
+        WellKnownTrait::FnOnceTrait => {
+            let interner = db.interner();
+            let self_ty = proj
+                .substitution
+                .at(interner, 0)
+                .assert_ty_ref(interner)
+                .data(interner);
+            let trait_id = db.well_known_trait_id(well_known).unwrap();
+            fn_::add_fn_trait_program_clauses(db, builder, trait_id, self_ty, true);
+        }
+        _ => {}
+    }
+}
+
 /// 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_.rs

@@ -1,27 +1,46 @@
 use crate::clauses::ClauseBuilder;
 use crate::infer::instantiate::IntoBindersAndValue;
+use crate::rust_ir::WellKnownTrait;
 use crate::{Interner, RustIrDatabase, TraitRef};
-use chalk_ir::{ApplicationTy, Binders, Substitution, Ty, TyData, TypeName, VariableKinds};
+use chalk_ir::{
+    AliasTy, ApplicationTy, Binders, Normalize, ProjectionTy, Substitution, TraitId, Ty, TyData,
+    TypeName, VariableKinds,
+};
 
-// Handles clauses for FnOnce/FnMut/Fn
+/// Handles clauses for FnOnce/FnMut/Fn.
+/// If `assoc_output` is `true`, we push a clause of the form
+/// `Normalize(<fn(A) -> B as FnOnce<(A,)>>::Output -> B) :- Implemented(fn(A) -> B as FnOnce<(A,)>`
+///
+/// If `assoc_output` is `false`, we push a clause of the form
+/// `Implemented(fn(A) -> B as FnOnce<(A,)>)`
 pub fn add_fn_trait_program_clauses<I: Interner>(
     db: &dyn RustIrDatabase<I>,
     builder: &mut ClauseBuilder<'_, I>,
-    trait_ref: &TraitRef<I>,
+    trait_id: TraitId<I>,
     ty: &TyData<I>,
+    assoc_output: bool,
 ) {
     match ty {
         TyData::Function(fn_val) => {
             let interner = db.interner();
-            let (binders, sub) = fn_val.into_binders_and_value(interner);
+            let (binders, orig_sub) = fn_val.into_binders_and_value(interner);
+            // Take all of the arguments except for the last one, which
+            // represents the return type
+            let arg_sub = Substitution::from(
+                interner,
+                orig_sub.iter(interner).take(orig_sub.len(interner) - 1),
+            );
+            let fn_output_ty = orig_sub
+                .at(interner, orig_sub.len(interner) - 1)
+                .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(sub.len(interner)),
-                    substitution: sub.clone(),
+                    name: TypeName::Tuple(arg_sub.len(interner)),
+                    substitution: arg_sub.clone(),
                 }),
             );
 
@@ -31,16 +50,42 @@ pub fn add_fn_trait_program_clauses<I: Interner>(
             // 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: trait_ref.trait_id,
-                substitution: tupled_sub,
+                trait_id,
+                substitution: tupled_sub.clone(),
             };
 
             // Functions types come with a binder, which we push so
             // that the `TraitRef` properly references any bound lifetimes
             // (e.g. `for<'a> fn(&'a u8): FnOnce<(&'b u8)>`)
             let bound_ref = Binders::new(VariableKinds::from(interner, binders), new_trait_ref);
             builder.push_binders(&bound_ref, |this, inner_trait| {
-                this.push_fact(inner_trait);
+                if assoc_output {
+                    //The `Output` type is defined on the `FnOnceTrait`
+                    let fn_once = db.trait_datum(trait_id);
+                    assert_eq!(fn_once.well_known, Some(WellKnownTrait::FnOnceTrait));
+                    let assoc_types = &fn_once.associated_ty_ids;
+                    if assoc_types.len() != 1 {
+                        panic!(
+                            "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(),
+                    };
+
+                    this.push_clause(normalize, std::iter::once(inner_trait));
+                } else {
+                    this.push_fact(inner_trait);
+                }
             })
         }
         _ => {}

tests/test/functions.rs

@@ -1,11 +1,13 @@
 use super::*;
 
 #[test]
-fn function_implement_fn_once() {
+fn function_implement_fn_traits() {
     test! {
         program {
             #[lang(fn_once)]
-            trait FnOnce<Args> { }
+            trait FnOnce<Args> {
+                type Output;
+            }
 
             #[lang(fn_mut)]
             trait FnMut<Args> where Self: FnOnce<Args> { }
@@ -32,6 +34,42 @@ fn function_implement_fn_once() {
             "Unique; substitution [], lifetime constraints []"
         }
 
+        goal {
+            Normalize(<fn(u8) as FnOnce<(u8,)>>::Output -> ())
+        } yields {
+            "Unique; substitution [], lifetime constraints []"
+        }
+
+        goal {
+            Normalize(<fn(u8) -> bool as FnOnce<(u8,)>>::Output -> bool)
+        } yields {
+            "Unique; substitution [], lifetime constraints []"
+        }
+
+        goal {
+            Normalize(<fn(u8) -> bool as FnOnce<(u8,)>>::Output -> u8)
+        } yields {
+            "No possible solution"
+        }
+
+        goal {
+            forall<T, V> {
+                Normalize(<fn(u8, V) -> T as FnOnce<(u8, V)>>::Output -> V)
+            }
+        } yields {
+            "No possible solution"
+        }
+
+        goal {
+            forall<T, V> {
+                exists<U> {
+                    Normalize(<fn(u8, V) -> T as FnOnce<(u8, V)>>::Output -> U)
+                }
+            }
+        } yields {
+            "Unique; substitution [?0 := !1_0], lifetime constraints []"
+        }
+
         goal {
             fn(u8, u32): FnOnce<(u8,u32)>
         } yields {