Add Type::walk method

3 files changed, 132 insertions, 75 deletions

diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs
index 5137a16e6..d0a8199a6 100644
--- a/crates/ra_hir/src/code_model.rs
+++ b/crates/ra_hir/src/code_model.rs
@@ -26,8 +26,8 @@ use hir_ty::{
     autoderef,
     display::{HirDisplayError, HirFormatter},
     expr::ExprValidator,
-    method_resolution, ApplicationTy, Canonical, GenericPredicate, InEnvironment, OpaqueTyId,
-    Substs, TraitEnvironment, Ty, TyDefId, TypeCtor,
+    method_resolution, ApplicationTy, Canonical, InEnvironment, Substs, TraitEnvironment, TraitRef,
+    Ty, TyDefId, TypeCtor,
 };
 use ra_db::{CrateId, CrateName, Edition, FileId};
 use ra_prof::profile;
@@ -1375,6 +1375,18 @@ impl Type {
         Some(adt.into())
     }
 
+    pub fn as_dyn_trait(&self) -> Option<Trait> {
+        self.ty.value.dyn_trait().map(Into::into)
+    }
+
+    pub fn as_impl_trait(&self, db: &dyn HirDatabase) -> Option<Trait> {
+        self.ty.value.impl_trait_ref(db).map(|it| it.trait_.into())
+    }
+
+    pub fn as_associated_type_parent_trait(&self, db: &dyn HirDatabase) -> Option<Trait> {
+        self.ty.value.associated_type_parent_trait(db).map(Into::into)
+    }
+
     // FIXME: provide required accessors such that it becomes implementable from outside.
     pub fn is_equal_for_find_impls(&self, other: &Type) -> bool {
         match (&self.ty.value, &other.ty.value) {
@@ -1397,96 +1409,72 @@ impl Type {
         }
     }
 
-    /// Returns a flattened list of all ADTs and Traits mentioned in the type
-    pub fn flattened_type_items(&self, db: &dyn HirDatabase) -> Vec<ModuleDef> {
-        fn push_new_item(item: ModuleDef, acc: &mut Vec<ModuleDef>) {
-            if !acc.contains(&item) {
-                acc.push(item);
+    pub fn walk(&self, db: &dyn HirDatabase, mut cb: impl FnMut(Type)) {
+        // TypeWalk::walk does not preserve items order!
+        fn walk_substs(db: &dyn HirDatabase, substs: &Substs, cb: &mut impl FnMut(Type)) {
+            for ty in substs.iter() {
+                walk_ty(db, ty, cb);
             }
         }
 
-        fn push_bounds(
+        fn walk_trait(
             db: &dyn HirDatabase,
-            predicates: &[GenericPredicate],
-            acc: &mut Vec<ModuleDef>,
+            ty: Ty,
+            trait_ref: &TraitRef,
+            cb: &mut impl FnMut(Type),
         ) {
-            for p in predicates.iter() {
-                match p {
-                    GenericPredicate::Implemented(trait_ref) => {
-                        push_new_item(Trait::from(trait_ref.trait_).into(), acc);
-                        walk_substs(db, &trait_ref.substs, acc);
-                    }
-                    GenericPredicate::Projection(_) => {}
-                    GenericPredicate::Error => (),
-                }
-            }
+            let def_db: &dyn DefDatabase = db.upcast();
+            let resolver = trait_ref.trait_.resolver(def_db);
+            let krate = trait_ref.trait_.lookup(def_db).container.module(def_db).krate;
+            cb(Type::new_with_resolver_inner(db, krate, &resolver, ty));
+            walk_substs(db, &trait_ref.substs, cb);
         }
 
-        // TypeWalk::walk does not preserve items order!
-        fn walk_substs(db: &dyn HirDatabase, substs: &Substs, acc: &mut Vec<ModuleDef>) {
-            for ty in substs.iter() {
-                walk_type(db, ty, acc);
-            }
-        }
-
-        fn walk_type(db: &dyn HirDatabase, ty: &Ty, acc: &mut Vec<ModuleDef>) {
-            match ty.strip_references() {
-                Ty::Apply(ApplicationTy { ctor, parameters, .. }) => {
+        fn walk_ty(db: &dyn HirDatabase, ty: &Ty, cb: &mut impl FnMut(Type)) {
+            let def_db: &dyn DefDatabase = db.upcast();
+            let ty = ty.strip_references();
+            match ty {
+                Ty::Apply(ApplicationTy { ctor, parameters }) => {
                     match ctor {
-                        TypeCtor::Adt(adt_id) => push_new_item(Adt::from(*adt_id).into(), acc),
-                        TypeCtor::AssociatedType(type_alias_id) => {
-                            let trait_id = match type_alias_id.lookup(db.upcast()).container {
-                                AssocContainerId::TraitId(it) => it,
-                                _ => panic!("not an associated type"),
-                            };
-
-                            push_new_item(Trait::from(trait_id).into(), acc);
+                        TypeCtor::Adt(adt) => {
+                            cb(Type::from_def(db, adt.module(def_db).krate, *adt));
+                        }
+                        TypeCtor::AssociatedType(_) => {
+                            if let Some(trait_id) = ty.associated_type_parent_trait(db) {
+                                let resolver = trait_id.resolver(def_db);
+                                let krate = trait_id.lookup(def_db).container.module(def_db).krate;
+                                cb(Type::new_with_resolver_inner(db, krate, &resolver, ty.clone()));
+                            }
                         }
                         _ => (),
                     }
+
                     // adt params, tuples, etc...
-                    walk_substs(db, parameters, acc);
+                    walk_substs(db, parameters, cb);
                 }
-                Ty::Dyn(predicates) => {
-                    push_bounds(db, predicates, acc);
+                Ty::Opaque(opaque_ty) => {
+                    if let Some(trait_ref) = ty.impl_trait_ref(db) {
+                        walk_trait(db, ty.clone(), &trait_ref, cb);
+                    }
+
+                    walk_substs(db, &opaque_ty.parameters, cb);
                 }
-                Ty::Placeholder(id) => {
-                    let generic_params = db.generic_params(id.parent);
-                    let param_data = &generic_params.types[id.local_id];
-                    match param_data.provenance {
-                        hir_def::generics::TypeParamProvenance::ArgumentImplTrait => {
-                            let predicates: Vec<_> = db
-                                .generic_predicates_for_param(*id)
-                                .into_iter()
-                                .map(|pred| pred.value.clone())
-                                .collect();
-                            push_bounds(db, &predicates, acc);
-                        }
-                        _ => (),
+                Ty::Placeholder(_) => {
+                    if let Some(trait_ref) = ty.impl_trait_ref(db) {
+                        walk_trait(db, ty.clone(), &trait_ref, cb);
                     }
                 }
-                Ty::Opaque(opaque_ty) => {
-                    let bounds = match opaque_ty.opaque_ty_id {
-                        OpaqueTyId::ReturnTypeImplTrait(func, idx) => {
-                            let datas = db
-                                .return_type_impl_traits(func)
-                                .expect("impl trait id without data");
-                            let data = (*datas)
-                                .as_ref()
-                                .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
-                            data.clone().subst(&opaque_ty.parameters)
-                        }
-                    };
-                    push_bounds(db, &bounds.value, acc);
-                    walk_substs(db, &opaque_ty.parameters, acc);
+                Ty::Dyn(_) => {
+                    if let Some(trait_ref) = ty.dyn_trait_ref() {
+                        walk_trait(db, ty.clone(), trait_ref, cb);
+                    }
                 }
+
                 _ => (),
             }
         }
 
-        let mut res: Vec<ModuleDef> = Vec::new(); // not a Set to preserve the order
-        walk_type(db, &self.ty.value, &mut res);
-        res
+        walk_ty(db, &self.ty.value, &mut cb);
     }
 }
 
diff --git a/crates/ra_hir_ty/src/lib.rs b/crates/ra_hir_ty/src/lib.rs
index 9d4d6aaa4..25ab9d872 100644
--- a/crates/ra_hir_ty/src/lib.rs
+++ b/crates/ra_hir_ty/src/lib.rs
@@ -877,6 +877,58 @@ impl Ty {
             _ => None,
         }
     }
+
+    pub fn impl_trait_ref(&self, db: &dyn HirDatabase) -> Option<TraitRef> {
+        match self {
+            Ty::Opaque(opaque_ty) => {
+                let predicates = match opaque_ty.opaque_ty_id {
+                    OpaqueTyId::ReturnTypeImplTrait(func, idx) => {
+                        db.return_type_impl_traits(func).map(|it| {
+                            let data = (*it)
+                                .as_ref()
+                                .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
+                            data.clone().subst(&opaque_ty.parameters)
+                        })
+                    }
+                };
+
+                predicates.and_then(|it| {
+                    it.value.iter().find_map(|pred| match pred {
+                        GenericPredicate::Implemented(tr) => Some(tr.clone()),
+                        _ => None,
+                    })
+                })
+            }
+            Ty::Placeholder(id) => {
+                let generic_params = db.generic_params(id.parent);
+                let param_data = &generic_params.types[id.local_id];
+                match param_data.provenance {
+                    hir_def::generics::TypeParamProvenance::ArgumentImplTrait => db
+                        .generic_predicates_for_param(*id)
+                        .into_iter()
+                        .map(|pred| pred.value.clone())
+                        .find_map(|pred| match pred {
+                            GenericPredicate::Implemented(tr) => Some(tr.clone()),
+                            _ => None,
+                        }),
+                    _ => None,
+                }
+            }
+            _ => None,
+        }
+    }
+
+    pub fn associated_type_parent_trait(&self, db: &dyn HirDatabase) -> Option<TraitId> {
+        match self {
+            Ty::Apply(ApplicationTy { ctor: TypeCtor::AssociatedType(type_alias_id), .. }) => {
+                match type_alias_id.lookup(db.upcast()).container {
+                    AssocContainerId::TraitId(trait_id) => Some(trait_id),
+                    _ => None,
+                }
+            }
+            _ => None,
+        }
+    }
 }
 
 /// This allows walking structures that contain types to do something with those
diff --git a/crates/ra_ide/src/hover.rs b/crates/ra_ide/src/hover.rs
index 045713519..c2909e200 100644
--- a/crates/ra_ide/src/hover.rs
+++ b/crates/ra_ide/src/hover.rs
@@ -236,9 +236,26 @@ fn runnable_action(
 fn goto_type_action(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
     match def {
         Definition::Local(it) => {
-            let targets = it
-                .ty(db)
-                .flattened_type_items(db)
+            let mut targets: Vec<ModuleDef> = Vec::new();
+            let mut push_new_def = |item: ModuleDef| {
+                if !targets.contains(&item) {
+                    targets.push(item);
+                }
+            };
+
+            it.ty(db).walk(db, |t| {
+                if let Some(adt) = t.as_adt() {
+                    push_new_def(adt.into());
+                } else if let Some(trait_) = t.as_dyn_trait() {
+                    push_new_def(trait_.into());
+                } else if let Some(trait_) = t.as_impl_trait(db) {
+                    push_new_def(trait_.into());
+                } else if let Some(trait_) = t.as_associated_type_parent_trait(db) {
+                    push_new_def(trait_.into());
+                }
+            });
+
+            let targets = targets
                 .into_iter()
                 .filter_map(|it| {
                     Some(HoverGotoTypeData {
@@ -246,7 +263,7 @@ fn goto_type_action(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
                         nav: it.try_to_nav(db)?,
                     })
                 })
-                .collect_vec();
+                .collect();
 
             Some(HoverAction::GoToType(targets))
         }