Add support for associated type bindings (`where Trait<Type = X>`)

4 files changed, 135 insertions, 46 deletions

diff --git a/crates/ra_hir/src/ty/lower.rs b/crates/ra_hir/src/ty/lower.rs
index 47d161277..0011c06b4 100644
--- a/crates/ra_hir/src/ty/lower.rs
+++ b/crates/ra_hir/src/ty/lower.rs
@@ -8,7 +8,9 @@
 use std::iter;
 use std::sync::Arc;
 
-use super::{FnSig, GenericPredicate, ProjectionTy, Substs, TraitRef, Ty, TypeCtor};
+use super::{
+    FnSig, GenericPredicate, ProjectionPredicate, ProjectionTy, Substs, TraitRef, Ty, TypeCtor,
+};
 use crate::{
     adt::VariantDef,
     generics::HasGenericParams,
@@ -62,7 +64,9 @@ impl Ty {
                 let self_ty = Ty::Bound(0);
                 let predicates = bounds
                     .iter()
-                    .map(|b| GenericPredicate::from_type_bound(db, resolver, b, self_ty.clone()))
+                    .flat_map(|b| {
+                        GenericPredicate::from_type_bound(db, resolver, b, self_ty.clone())
+                    })
                     .collect::<Vec<_>>();
                 Ty::Dyn(predicates.into())
             }
@@ -70,7 +74,9 @@ impl Ty {
                 let self_ty = Ty::Bound(0);
                 let predicates = bounds
                     .iter()
-                    .map(|b| GenericPredicate::from_type_bound(db, resolver, b, self_ty.clone()))
+                    .flat_map(|b| {
+                        GenericPredicate::from_type_bound(db, resolver, b, self_ty.clone())
+                    })
                     .collect::<Vec<_>>();
                 Ty::Opaque(predicates.into())
             }
@@ -326,15 +332,6 @@ impl TraitRef {
         TraitRef { trait_, substs }
     }
 
-    pub(crate) fn from_where_predicate(
-        db: &impl HirDatabase,
-        resolver: &Resolver,
-        pred: &WherePredicate,
-    ) -> Option<TraitRef> {
-        let self_ty = Ty::from_hir(db, resolver, &pred.type_ref);
-        TraitRef::from_type_bound(db, resolver, &pred.bound, self_ty)
-    }
-
     pub(crate) fn from_type_bound(
         db: &impl HirDatabase,
         resolver: &Resolver,
@@ -349,26 +346,58 @@ impl TraitRef {
 }
 
 impl GenericPredicate {
-    pub(crate) fn from_where_predicate(
-        db: &impl HirDatabase,
-        resolver: &Resolver,
-        where_predicate: &WherePredicate,
-    ) -> GenericPredicate {
-        TraitRef::from_where_predicate(db, &resolver, where_predicate)
-            .map_or(GenericPredicate::Error, GenericPredicate::Implemented)
+    pub(crate) fn from_where_predicate<'a>(
+        db: &'a impl HirDatabase,
+        resolver: &'a Resolver,
+        where_predicate: &'a WherePredicate,
+    ) -> impl Iterator<Item = GenericPredicate> + 'a {
+        let self_ty = Ty::from_hir(db, resolver, &where_predicate.type_ref);
+        GenericPredicate::from_type_bound(db, resolver, &where_predicate.bound, self_ty)
     }
 
-    pub(crate) fn from_type_bound(
-        db: &impl HirDatabase,
-        resolver: &Resolver,
-        bound: &TypeBound,
+    pub(crate) fn from_type_bound<'a>(
+        db: &'a impl HirDatabase,
+        resolver: &'a Resolver,
+        bound: &'a TypeBound,
         self_ty: Ty,
-    ) -> GenericPredicate {
-        TraitRef::from_type_bound(db, &resolver, bound, self_ty)
-            .map_or(GenericPredicate::Error, GenericPredicate::Implemented)
+    ) -> impl Iterator<Item = GenericPredicate> + 'a {
+        let trait_ref = TraitRef::from_type_bound(db, &resolver, bound, self_ty);
+        iter::once(trait_ref.clone().map_or(GenericPredicate::Error, GenericPredicate::Implemented))
+            .chain(
+                trait_ref.into_iter().flat_map(move |tr| {
+                    assoc_type_bindings_from_type_bound(db, resolver, bound, tr)
+                }),
+            )
     }
 }
 
+fn assoc_type_bindings_from_type_bound<'a>(
+    db: &'a impl HirDatabase,
+    resolver: &'a Resolver,
+    bound: &'a TypeBound,
+    trait_ref: TraitRef,
+) -> impl Iterator<Item = GenericPredicate> + 'a {
+    let last_segment = match bound {
+        TypeBound::Path(path) => path.segments.last(),
+        TypeBound::Error => None,
+    };
+    last_segment
+        .into_iter()
+        .flat_map(|segment| segment.args_and_bindings.iter())
+        .flat_map(|args_and_bindings| args_and_bindings.bindings.iter())
+        .map(move |(name, type_ref)| {
+            let associated_ty = match trait_ref.trait_.associated_type_by_name(db, name.clone()) {
+                None => return GenericPredicate::Error,
+                Some(t) => t,
+            };
+            let projection_ty =
+                ProjectionTy { associated_ty, parameters: trait_ref.substs.clone() };
+            let ty = Ty::from_hir(db, resolver, type_ref);
+            let projection_predicate = ProjectionPredicate { projection_ty, ty };
+            GenericPredicate::Projection(projection_predicate)
+        })
+}
+
 /// Build the declared type of an item. This depends on the namespace; e.g. for
 /// `struct Foo(usize)`, we have two types: The type of the struct itself, and
 /// the constructor function `(usize) -> Foo` which lives in the values
@@ -425,7 +454,7 @@ pub(crate) fn trait_env(
 ) -> Arc<super::TraitEnvironment> {
     let predicates = resolver
         .where_predicates_in_scope()
-        .map(|pred| GenericPredicate::from_where_predicate(db, &resolver, pred))
+        .flat_map(|pred| GenericPredicate::from_where_predicate(db, &resolver, pred))
         .collect::<Vec<_>>();
 
     Arc::new(super::TraitEnvironment { predicates })
@@ -439,7 +468,7 @@ pub(crate) fn generic_predicates_query(
     let resolver = def.resolver(db);
     let predicates = resolver
         .where_predicates_in_scope()
-        .map(|pred| GenericPredicate::from_where_predicate(db, &resolver, pred))
+        .flat_map(|pred| GenericPredicate::from_where_predicate(db, &resolver, pred))
         .collect::<Vec<_>>();
     predicates.into()
 }
diff --git a/crates/ra_hir/src/ty/tests.rs b/crates/ra_hir/src/ty/tests.rs
index 42bcae157..3e743ef58 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -3586,37 +3586,63 @@ fn test<T: Trait<Type = u32>>(x: T, y: impl Trait<Type = i64>) {
     [166; 169) '{t}': T
     [167; 168) 't': T
     [257; 258) 'x': T
-    [263; 264) 'y': impl Trait
+    [263; 264) 'y': impl Trait + 
     [290; 398) '{     ...r>); }': ()
     [296; 299) 'get': fn get<T>(T) -> <T as Trait>::Type
     [296; 302) 'get(x)': {unknown}
     [300; 301) 'x': T
-    [308; 312) 'get2': fn get2<{unknown}, T>(T) -> U
+    [308; 312) 'get2': fn get2<{unknown}, S<{unknown}>>(T) -> U
     [308; 315) 'get2(x)': {unknown}
     [313; 314) 'x': T
-    [321; 324) 'get': fn get<impl Trait>(T) -> <T as Trait>::Type
+    [321; 324) 'get': fn get<impl Trait + >(T) -> <T as Trait>::Type
     [321; 327) 'get(y)': {unknown}
-    [325; 326) 'y': impl Trait
-    [333; 337) 'get2': fn get2<{unknown}, impl Trait>(T) -> U
+    [325; 326) 'y': impl Trait + 
+    [333; 337) 'get2': fn get2<{unknown}, S<{unknown}>>(T) -> U
     [333; 340) 'get2(y)': {unknown}
-    [338; 339) 'y': impl Trait
-    [346; 349) 'get': fn get<S<{unknown}>>(T) -> <T as Trait>::Type
-    [346; 357) 'get(set(S))': {unknown}
-    [350; 353) 'set': fn set<S<{unknown}>>(T) -> T
-    [350; 356) 'set(S)': S<{unknown}>
-    [354; 355) 'S': S<{unknown}>
-    [363; 367) 'get2': fn get2<{unknown}, S<{unknown}>>(T) -> U
-    [363; 375) 'get2(set(S))': {unknown}
-    [368; 371) 'set': fn set<S<{unknown}>>(T) -> T
-    [368; 374) 'set(S)': S<{unknown}>
-    [372; 373) 'S': S<{unknown}>
-    [381; 385) 'get2': fn get2<{unknown}, S<str>>(T) -> U
-    [381; 395) 'get2(S::<str>)': {unknown}
+    [338; 339) 'y': impl Trait + 
+    [346; 349) 'get': fn get<S<u64>>(T) -> <T as Trait>::Type
+    [346; 357) 'get(set(S))': u64
+    [350; 353) 'set': fn set<S<u64>>(T) -> T
+    [350; 356) 'set(S)': S<u64>
+    [354; 355) 'S': S<u64>
+    [363; 367) 'get2': fn get2<u64, S<u64>>(T) -> U
+    [363; 375) 'get2(set(S))': u64
+    [368; 371) 'set': fn set<S<u64>>(T) -> T
+    [368; 374) 'set(S)': S<u64>
+    [372; 373) 'S': S<u64>
+    [381; 385) 'get2': fn get2<str, S<str>>(T) -> U
+    [381; 395) 'get2(S::<str>)': str
     [386; 394) 'S::<str>': S<str>
     "###
     );
 }
 
+#[test]
+fn projection_eq_within_chalk() {
+    // std::env::set_var("CHALK_DEBUG", "1");
+    assert_snapshot!(
+        infer(r#"
+trait Trait1 {
+    type Type;
+}
+trait Trait2<T> {
+    fn foo(self) -> T;
+}
+impl<T, U> Trait2<T> for U where U: Trait1<Type = T> {}
+
+fn test<T: Trait1<Type = u32>>(x: T) {
+    x.foo();
+}
+"#),
+        @r###"
+    [62; 66) 'self': Self
+    [164; 165) 'x': T
+    [170; 186) '{     ...o(); }': ()
+    [176; 177) 'x': T
+    [176; 183) 'x.foo()': {unknown}
+    "###
+    );
+}
 fn type_at_pos(db: &MockDatabase, pos: FilePosition) -> String {
     let file = db.parse(pos.file_id).ok().unwrap();
     let expr = algo::find_node_at_offset::<ast::Expr>(file.syntax(), pos.offset).unwrap();
diff --git a/crates/ra_hir/src/ty/traits.rs b/crates/ra_hir/src/ty/traits.rs
index b634f0b79..25316bc02 100644
--- a/crates/ra_hir/src/ty/traits.rs
+++ b/crates/ra_hir/src/ty/traits.rs
@@ -124,6 +124,9 @@ impl Obligation {
     pub fn from_predicate(predicate: GenericPredicate) -> Option<Obligation> {
         match predicate {
             GenericPredicate::Implemented(trait_ref) => Some(Obligation::Trait(trait_ref)),
+            GenericPredicate::Projection(projection_pred) => {
+                Some(Obligation::Projection(projection_pred))
+            }
             GenericPredicate::Error => None,
         }
     }
@@ -135,6 +138,30 @@ pub struct ProjectionPredicate {
     pub ty: Ty,
 }
 
+impl ProjectionPredicate {
+    pub fn subst(mut self, substs: &super::Substs) -> ProjectionPredicate {
+        self.walk_mut(&mut |ty| match ty {
+            Ty::Param { idx, .. } => {
+                if let Some(t) = substs.get(*idx as usize).cloned() {
+                    *ty = t;
+                }
+            }
+            _ => {}
+        });
+        self
+    }
+
+    pub fn walk(&self, f: &mut impl FnMut(&Ty)) {
+        self.projection_ty.walk(f);
+        self.ty.walk(f);
+    }
+
+    pub fn walk_mut(&mut self, f: &mut impl FnMut(&mut Ty)) {
+        self.projection_ty.walk_mut(f);
+        self.ty.walk_mut(f);
+    }
+}
+
 /// Solve a trait goal using Chalk.
 pub(crate) fn trait_solve_query(
     db: &impl HirDatabase,
diff --git a/crates/ra_hir/src/ty/traits/chalk.rs b/crates/ra_hir/src/ty/traits/chalk.rs
index 2ebc06135..3ab5b7cca 100644
--- a/crates/ra_hir/src/ty/traits/chalk.rs
+++ b/crates/ra_hir/src/ty/traits/chalk.rs
@@ -211,6 +211,13 @@ impl ToChalk for GenericPredicate {
             GenericPredicate::Implemented(trait_ref) => {
                 make_binders(chalk_ir::WhereClause::Implemented(trait_ref.to_chalk(db)), 0)
             }
+            GenericPredicate::Projection(projection_pred) => make_binders(
+                chalk_ir::WhereClause::ProjectionEq(chalk_ir::ProjectionEq {
+                    projection: projection_pred.projection_ty.to_chalk(db),
+                    ty: projection_pred.ty.to_chalk(db),
+                }),
+                0,
+            ),
             GenericPredicate::Error => {
                 let impossible_trait_ref = chalk_ir::TraitRef {
                     trait_id: UNKNOWN_TRAIT,