Get trait assoc item resolution mostly working

3 files changed, 261 insertions, 16 deletions

diff --git a/crates/ra_hir/src/ty.rs b/crates/ra_hir/src/ty.rs
index d2bfcdc7d..e39f06e68 100644
--- a/crates/ra_hir/src/ty.rs
+++ b/crates/ra_hir/src/ty.rs
@@ -385,13 +385,21 @@ impl SubstsBuilder {
         self.param_count - self.vec.len()
     }
 
-    pub fn fill_with_bound_vars(mut self, starting_from: u32) -> Self {
-        self.vec.extend((starting_from..starting_from + self.remaining() as u32).map(Ty::Bound));
-        self
+    pub fn fill_with_bound_vars(self, starting_from: u32) -> Self {
+        self.fill((starting_from..).map(Ty::Bound))
+    }
+
+    pub fn fill_with_params(self) -> Self {
+        let start = self.vec.len() as u32;
+        self.fill((start..).map(|idx| Ty::Param { idx, name: Name::missing() }))
+    }
+
+    pub fn fill_with_unknown(self) -> Self {
+        self.fill(iter::repeat(Ty::Unknown))
     }
 
-    pub fn fill_with_unknown(mut self) -> Self {
-        self.vec.extend(iter::repeat(Ty::Unknown).take(self.remaining()));
+    pub fn fill(mut self, filler: impl Iterator<Item = Ty>) -> Self {
+        self.vec.extend(filler.take(self.remaining()));
         self
     }
 
diff --git a/crates/ra_hir/src/ty/infer/path.rs b/crates/ra_hir/src/ty/infer/path.rs
index 77aa35ce1..885588174 100644
--- a/crates/ra_hir/src/ty/infer/path.rs
+++ b/crates/ra_hir/src/ty/infer/path.rs
@@ -6,9 +6,11 @@ use super::{ExprOrPatId, InferenceContext, TraitRef};
 use crate::{
     db::HirDatabase,
     resolve::{ResolveValueResult, Resolver, TypeNs, ValueNs},
+    ty::{lower, traits::TraitEnvironment, Canonical},
     ty::{Substs, Ty, TypableDef, TypeWalk},
-    AssocItem, HasGenericParams, Namespace, Path,
+    AssocItem, HasGenericParams, Name, Namespace, Path, Trait,
 };
+use std::sync::Arc;
 
 impl<'a, D: HirDatabase> InferenceContext<'a, D> {
     pub(super) fn infer_path(
@@ -39,7 +41,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
             let ty = Ty::from_type_relative_path(self.db, resolver, ty, remaining_segments_for_ty);
             self.resolve_ty_assoc_item(
                 ty,
-                path.segments.last().expect("path had at least one segment"),
+                &path.segments.last().expect("path had at least one segment").name,
                 id,
             )?
         } else {
@@ -125,7 +127,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 let segment =
                     remaining_segments.last().expect("there should be at least one segment here");
 
-                self.resolve_ty_assoc_item(ty, segment, id)
+                self.resolve_ty_assoc_item(ty, &segment.name, id)
             }
         }
     }
@@ -162,7 +164,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         };
         let substs = Substs::build_for_def(self.db, item)
             .use_parent_substs(&trait_ref.substs)
-            .fill_with_unknown()
+            .fill_with_params()
             .build();
 
         self.write_assoc_resolution(id, item);
@@ -172,20 +174,29 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
     fn resolve_ty_assoc_item(
         &mut self,
         ty: Ty,
-        segment: &PathSegment,
+        name: &Name,
         id: ExprOrPatId,
     ) -> Option<(ValueNs, Option<Substs>)> {
         if let Ty::Unknown = ty {
             return None;
         }
 
+        self.find_inherent_assoc_candidate(ty.clone(), name, id)
+            .or_else(|| self.find_trait_assoc_candidate(ty.clone(), name, id))
+    }
+
+    fn find_inherent_assoc_candidate(
+        &mut self,
+        ty: Ty,
+        name: &Name,
+        id: ExprOrPatId,
+    ) -> Option<(ValueNs, Option<Substs>)> {
         let krate = self.resolver.krate()?;
 
         // Find impl
-        // FIXME: consider trait candidates
         let item = ty.clone().iterate_impl_items(self.db, krate, |item| match item {
             AssocItem::Function(func) => {
-                if segment.name == func.name(self.db) {
+                if *name == func.name(self.db) {
                     Some(AssocItem::Function(func))
                 } else {
                     None
@@ -193,7 +204,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
             }
 
             AssocItem::Const(konst) => {
-                if konst.name(self.db).map_or(false, |n| n == segment.name) {
+                if konst.name(self.db).map_or(false, |n| n == *name) {
                     Some(AssocItem::Const(konst))
                 } else {
                     None
@@ -212,6 +223,65 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         Some((def, substs))
     }
 
+    fn find_trait_assoc_candidate(
+        &mut self,
+        ty: Ty,
+        name: &Name,
+        _id: ExprOrPatId,
+    ) -> Option<(ValueNs, Option<Substs>)> {
+        let krate = self.resolver.krate()?;
+
+        let canonical_ty = self.canonicalizer().canonicalize_ty(ty.clone());
+
+        let env = lower::trait_env(self.db, &self.resolver);
+        // if we have `T: Trait` in the param env, the trait doesn't need to be in scope
+        let traits_from_env = env
+            .trait_predicates_for_self_ty(&ty)
+            .map(|tr| tr.trait_)
+            .flat_map(|t| t.all_super_traits(self.db));
+        let traits = traits_from_env.chain(self.resolver.traits_in_scope(self.db));
+
+        'traits: for t in traits {
+            let data = t.trait_data(self.db);
+            let mut known_implemented = false;
+            for item in data.items() {
+                if let AssocItem::Function(f) = *item {
+                    if f.name(self.db) == *name {
+                        if !known_implemented {
+                            let goal = generic_implements_goal(
+                                self.db,
+                                env.clone(),
+                                t,
+                                canonical_ty.value.clone(),
+                            );
+                            if self.db.trait_solve(krate, goal).is_none() {
+                                continue 'traits;
+                            }
+                        }
+                        known_implemented = true;
+
+                        // we're picking this method
+                        let trait_substs = Substs::build_for_def(self.db, t)
+                            .push(ty.clone())
+                            .fill(std::iter::repeat_with(|| self.new_type_var()))
+                            .build();
+                        let substs = Substs::build_for_def(self.db, f)
+                            .use_parent_substs(&trait_substs)
+                            .fill_with_params()
+                            .build();
+                        self.obligations.push(super::Obligation::Trait(TraitRef {
+                            trait_: t,
+                            substs: trait_substs,
+                        }));
+                        return Some((ValueNs::Function(f), Some(substs)));
+                    }
+                }
+            }
+        }
+
+        None
+    }
+
     fn find_self_types(&self, def: &ValueNs, actual_def_ty: Ty) -> Option<Substs> {
         if let ValueNs::Function(func) = def {
             // We only do the infer if parent has generic params
@@ -242,3 +312,23 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         }
     }
 }
+
+// TODO remove duplication
+/// This creates Substs for a trait with the given Self type and type variables
+/// for all other parameters, to query Chalk with it.
+fn generic_implements_goal(
+    db: &impl HirDatabase,
+    env: Arc<TraitEnvironment>,
+    trait_: Trait,
+    self_ty: Canonical<Ty>,
+) -> Canonical<super::InEnvironment<super::Obligation>> {
+    let num_vars = self_ty.num_vars;
+    let substs = super::Substs::build_for_def(db, trait_)
+        .push(self_ty.value)
+        .fill_with_bound_vars(num_vars as u32)
+        .build();
+    let num_vars = substs.len() - 1 + self_ty.num_vars;
+    let trait_ref = TraitRef { trait_, substs };
+    let obligation = super::Obligation::Trait(trait_ref);
+    Canonical { num_vars, value: super::InEnvironment::new(env, obligation) }
+}
diff --git a/crates/ra_hir/src/ty/tests.rs b/crates/ra_hir/src/ty/tests.rs
index c12326643..7183b205c 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -2782,9 +2782,9 @@ fn test() {
     [97; 99) 's1': S
     [105; 121) 'Defaul...efault': fn default<S>() -> Self
     [105; 123) 'Defaul...ault()': S
-    [133; 135) 's2': {unknown}
-    [138; 148) 'S::default': {unknown}
-    [138; 150) 'S::default()': {unknown}
+    [133; 135) 's2': S
+    [138; 148) 'S::default': fn default<S>() -> Self
+    [138; 150) 'S::default()': S
     [160; 162) 's3': S
     [165; 188) '<S as ...efault': fn default<S>() -> Self
     [165; 190) '<S as ...ault()': S
@@ -2793,6 +2793,153 @@ fn test() {
 }
 
 #[test]
+fn infer_trait_assoc_method_generics_1() {
+    assert_snapshot!(
+        infer(r#"
+trait Trait<T> {
+    fn make() -> T;
+}
+struct S;
+impl Trait<u32> for S {}
+struct G<T>;
+impl<T> Trait<T> for G<T> {}
+fn test() {
+    let a = S::make();
+    let b = G::<u64>::make();
+    let c: f64 = G::make();
+}
+"#),
+        @r###"
+    [127; 211) '{     ...e(); }': ()
+    [137; 138) 'a': u32
+    [141; 148) 'S::make': fn make<S, u32>() -> T
+    [141; 150) 'S::make()': u32
+    [160; 161) 'b': u64
+    [164; 178) 'G::<u64>::make': fn make<G<u64>, u64>() -> T
+    [164; 180) 'G::<u6...make()': u64
+    [190; 191) 'c': f64
+    [199; 206) 'G::make': fn make<G<f64>, f64>() -> T
+    [199; 208) 'G::make()': f64
+    "###
+    );
+}
+
+#[test]
+fn infer_trait_assoc_method_generics_2() {
+    assert_snapshot!(
+        infer(r#"
+trait Trait<T> {
+    fn make<U>() -> (T, U);
+}
+struct S;
+impl Trait<u32> for S {}
+struct G<T>;
+impl<T> Trait<T> for G<T> {}
+fn test() {
+    let a = S::make::<i64>();
+    let b: (_, i64) = S::make();
+    let c = G::<u32>::make::<i64>();
+    let d: (u32, _) = G::make::<i64>();
+    let e: (u32, i64) = G::make();
+}
+"#),
+        @r###"
+    [135; 313) '{     ...e(); }': ()
+    [145; 146) 'a': (u32, i64)
+    [149; 163) 'S::make::<i64>': fn make<S, u32, i64>() -> (T, U)
+    [149; 165) 'S::mak...i64>()': (u32, i64)
+    [175; 176) 'b': (u32, i64)
+    [189; 196) 'S::make': fn make<S, u32, i64>() -> (T, U)
+    [189; 198) 'S::make()': (u32, i64)
+    [208; 209) 'c': (u32, i64)
+    [212; 233) 'G::<u3...:<i64>': fn make<G<u32>, u32, i64>() -> (T, U)
+    [212; 235) 'G::<u3...i64>()': (u32, i64)
+    [245; 246) 'd': (u32, i64)
+    [259; 273) 'G::make::<i64>': fn make<G<u32>, u32, i64>() -> (T, U)
+    [259; 275) 'G::mak...i64>()': (u32, i64)
+    [285; 286) 'e': (u32, i64)
+    [301; 308) 'G::make': fn make<G<u32>, u32, i64>() -> (T, U)
+    [301; 310) 'G::make()': (u32, i64)
+    "###
+    );
+}
+
+#[test]
+fn infer_trait_assoc_method_generics_3() {
+    assert_snapshot!(
+        infer(r#"
+trait Trait<T> {
+    fn make() -> (Self, T);
+}
+struct S<T>;
+impl Trait<i64> for S<i32> {}
+fn test() {
+    let a = S::make();
+}
+"#),
+        @r###"
+    [101; 127) '{     ...e(); }': ()
+    [111; 112) 'a': {unknown}
+    [115; 122) 'S::make': {unknown}
+    [115; 124) 'S::make()': {unknown}
+    "###
+    );
+}
+
+#[test]
+fn infer_trait_assoc_method_generics_4() {
+    assert_snapshot!(
+        infer(r#"
+trait Trait<T> {
+    fn make() -> (Self, T);
+}
+struct S<T>;
+impl Trait<i64> for S<u64> {}
+impl Trait<i32> for S<u32> {}
+fn test() {
+    let a: (Self<i64>, _) = S::make();
+    let b: (_, u32) = S::make();
+}
+"#),
+        @r###"
+    [131; 206) '{     ...e(); }': ()
+    [141; 142) 'a': ({unknown}, {unknown})
+    [161; 168) 'S::make': {unknown}
+    [161; 170) 'S::make()': ({unknown}, {unknown})
+    [180; 181) 'b': ({unknown}, u32)
+    [194; 201) 'S::make': {unknown}
+    [194; 203) 'S::make()': ({unknown}, u32)
+    "###
+    );
+}
+
+#[test]
+fn infer_trait_assoc_method_generics_5() {
+    assert_snapshot!(
+        infer(r#"
+trait Trait<T> {
+    fn make<U>() -> (Self, T, U);
+}
+struct S<T>;
+impl Trait<i64> for S<u64> {}
+fn test() {
+    let a = <S as Trait<i64>>::make::<u8>();
+    let b: (S<u64>, _, _) = Trait::<i64>::make::<u8>();
+}
+"#),
+        @r###"
+    [107; 211) '{     ...>(); }': ()
+    [117; 118) 'a': (S<u64>, i64, u8)
+    [121; 150) '<S as ...::<u8>': fn make<S<u64>, i64, u8>() -> (Self, T, U)
+    [121; 152) '<S as ...<u8>()': (S<u64>, i64, u8)
+    [162; 163) 'b': (S<u64>, i64, u8)
+    [182; 206) 'Trait:...::<u8>': fn make<S<u64>, i64, u8>() -> (Self, T, U)
+    [182; 208) 'Trait:...<u8>()': (S<u64>, i64, u8)
+    "###
+    );
+}
+
+#[test]
 fn infer_from_bound_1() {
     assert_snapshot!(
         infer(r#"