1 files changed, 218 insertions, 152 deletions

diff --git a/crates/ra_hir_def/src/import_map.rs b/crates/ra_hir_def/src/import_map.rs
index 68e20d06b..9e4c30b1a 100644
--- a/crates/ra_hir_def/src/import_map.rs
+++ b/crates/ra_hir_def/src/import_map.rs
@@ -5,14 +5,16 @@ use std::{cmp::Ordering, fmt, hash::BuildHasherDefault, sync::Arc};
 use fst::{self, Streamer};
 use indexmap::{map::Entry, IndexMap};
 use ra_db::CrateId;
-use rustc_hash::FxHasher;
+use ra_syntax::SmolStr;
+use rustc_hash::{FxHashMap, FxHasher};
+use smallvec::SmallVec;
 
 use crate::{
     db::DefDatabase,
     item_scope::ItemInNs,
     path::{ModPath, PathKind},
     visibility::Visibility,
-    ModuleDefId, ModuleId,
+    AssocItemId, ModuleDefId, ModuleId, TraitId,
 };
 
 type FxIndexMap<K, V> = IndexMap<K, V, BuildHasherDefault<FxHasher>>;
@@ -34,6 +36,7 @@ pub struct ImportInfo {
 ///
 /// Note that all paths are relative to the containing crate's root, so the crate name still needs
 /// to be prepended to the `ModPath` before the path is valid.
+#[derive(Default)]
 pub struct ImportMap {
     map: FxIndexMap<ItemInNs, ImportInfo>,
 
@@ -45,13 +48,17 @@ pub struct ImportMap {
     /// the index of the first one.
     importables: Vec<ItemInNs>,
     fst: fst::Map<Vec<u8>>,
+
+    /// Maps names of associated items to the item's ID. Only includes items whose defining trait is
+    /// exported.
+    assoc_map: FxHashMap<SmolStr, SmallVec<[AssocItemId; 1]>>,
 }
 
 impl ImportMap {
     pub fn import_map_query(db: &dyn DefDatabase, krate: CrateId) -> Arc<Self> {
         let _p = ra_prof::profile("import_map_query");
         let def_map = db.crate_def_map(krate);
-        let mut import_map = FxIndexMap::with_capacity_and_hasher(64, Default::default());
+        let mut import_map = Self::default();
 
         // We look only into modules that are public(ly reexported), starting with the crate root.
         let empty = ModPath { kind: PathKind::Plain, segments: vec![] };
@@ -85,7 +92,7 @@ impl ImportMap {
 
                 for item in per_ns.iter_items() {
                     let path = mk_path();
-                    match import_map.entry(item) {
+                    match import_map.map.entry(item) {
                         Entry::Vacant(entry) => {
                             entry.insert(ImportInfo { path, container: module });
                         }
@@ -105,11 +112,16 @@ impl ImportMap {
                     if let Some(ModuleDefId::ModuleId(mod_id)) = item.as_module_def_id() {
                         worklist.push((mod_id, mk_path()));
                     }
+
+                    // If we've added a path to a trait, add the trait's methods to the method map.
+                    if let Some(ModuleDefId::TraitId(tr)) = item.as_module_def_id() {
+                        import_map.collect_trait_methods(db, tr);
+                    }
                 }
             }
         }
 
-        let mut importables = import_map.iter().collect::<Vec<_>>();
+        let mut importables = import_map.map.iter().collect::<Vec<_>>();
 
         importables.sort_by(cmp);
 
@@ -133,10 +145,10 @@ impl ImportMap {
             builder.insert(key, start as u64).unwrap();
         }
 
-        let fst = fst::Map::new(builder.into_inner().unwrap()).unwrap();
-        let importables = importables.iter().map(|(item, _)| **item).collect();
+        import_map.fst = fst::Map::new(builder.into_inner().unwrap()).unwrap();
+        import_map.importables = importables.iter().map(|(item, _)| **item).collect();
 
-        Arc::new(Self { map: import_map, fst, importables })
+        Arc::new(import_map)
     }
 
     /// Returns the `ModPath` needed to import/mention `item`, relative to this crate's root.
@@ -147,6 +159,13 @@ impl ImportMap {
     pub fn import_info_for(&self, item: ItemInNs) -> Option<&ImportInfo> {
         self.map.get(&item)
     }
+
+    fn collect_trait_methods(&mut self, db: &dyn DefDatabase, tr: TraitId) {
+        let data = db.trait_data(tr);
+        for (name, item) in data.items.iter() {
+            self.assoc_map.entry(name.to_string().into()).or_default().push(*item);
+        }
+    }
 }
 
 impl PartialEq for ImportMap {
@@ -290,37 +309,32 @@ pub fn search_dependencies<'a>(
         }
     }
 
+    // Add all exported associated items whose names match the query (exactly).
+    for map in &import_maps {
+        if let Some(v) = map.assoc_map.get(&*query.query) {
+            res.extend(v.iter().map(|&assoc| {
+                ItemInNs::Types(match assoc {
+                    AssocItemId::FunctionId(it) => it.into(),
+                    AssocItemId::ConstId(it) => it.into(),
+                    AssocItemId::TypeAliasId(it) => it.into(),
+                })
+            }));
+        }
+    }
+
     res
 }
 
 #[cfg(test)]
 mod tests {
-    use super::*;
-    use crate::test_db::TestDB;
-    use insta::assert_snapshot;
-    use itertools::Itertools;
-    use ra_db::fixture::WithFixture;
-    use ra_db::{SourceDatabase, Upcast};
+    use expect::{expect, Expect};
+    use ra_db::{fixture::WithFixture, SourceDatabase, Upcast};
 
-    fn import_map(ra_fixture: &str) -> String {
-        let db = TestDB::with_files(ra_fixture);
-        let crate_graph = db.crate_graph();
+    use crate::{test_db::TestDB, AssocContainerId, Lookup};
 
-        let s = crate_graph
-            .iter()
-            .filter_map(|krate| {
-                let cdata = &crate_graph[krate];
-                let name = cdata.display_name.as_ref()?;
-
-                let map = db.import_map(krate);
-
-                Some(format!("{}:\n{:?}", name, map))
-            })
-            .join("\n");
-        s
-    }
+    use super::*;
 
-    fn search_dependencies_of(ra_fixture: &str, krate_name: &str, query: Query) -> String {
+    fn check_search(ra_fixture: &str, krate_name: &str, query: Query, expect: Expect) {
         let db = TestDB::with_files(ra_fixture);
         let crate_graph = db.crate_graph();
         let krate = crate_graph
@@ -331,7 +345,7 @@ mod tests {
             })
             .unwrap();
 
-        search_dependencies(db.upcast(), krate, query)
+        let actual = search_dependencies(db.upcast(), krate, query)
             .into_iter()
             .filter_map(|item| {
                 let mark = match item {
@@ -339,23 +353,67 @@ mod tests {
                     ItemInNs::Values(_) => "v",
                     ItemInNs::Macros(_) => "m",
                 };
+                let item = assoc_to_trait(&db, item);
                 item.krate(db.upcast()).map(|krate| {
                     let map = db.import_map(krate);
                     let path = map.path_of(item).unwrap();
                     format!(
-                        "{}::{} ({})",
+                        "{}::{} ({})\n",
                         crate_graph[krate].display_name.as_ref().unwrap(),
                         path,
                         mark
                     )
                 })
             })
-            .join("\n")
+            .collect::<String>();
+        expect.assert_eq(&actual)
+    }
+
+    fn assoc_to_trait(db: &dyn DefDatabase, item: ItemInNs) -> ItemInNs {
+        let assoc: AssocItemId = match item {
+            ItemInNs::Types(it) | ItemInNs::Values(it) => match it {
+                ModuleDefId::TypeAliasId(it) => it.into(),
+                ModuleDefId::FunctionId(it) => it.into(),
+                ModuleDefId::ConstId(it) => it.into(),
+                _ => return item,
+            },
+            _ => return item,
+        };
+
+        let container = match assoc {
+            AssocItemId::FunctionId(it) => it.lookup(db).container,
+            AssocItemId::ConstId(it) => it.lookup(db).container,
+            AssocItemId::TypeAliasId(it) => it.lookup(db).container,
+        };
+
+        match container {
+            AssocContainerId::TraitId(it) => ItemInNs::Types(it.into()),
+            _ => item,
+        }
+    }
+
+    fn check(ra_fixture: &str, expect: Expect) {
+        let db = TestDB::with_files(ra_fixture);
+        let crate_graph = db.crate_graph();
+
+        let actual = crate_graph
+            .iter()
+            .filter_map(|krate| {
+                let cdata = &crate_graph[krate];
+                let name = cdata.display_name.as_ref()?;
+
+                let map = db.import_map(krate);
+
+                Some(format!("{}:\n{:?}\n", name, map))
+            })
+            .collect::<String>();
+
+        expect.assert_eq(&actual)
     }
 
     #[test]
     fn smoke() {
-        let map = import_map(
+        check(
             r"
             //- /main.rs crate:main deps:lib
 
@@ -380,24 +438,23 @@ mod tests {
             pub struct Pub2; // t + v
             struct Priv;
         ",
+            expect![[r#"
+                main:
+                - publ1 (t)
+                - real_pu2 (t)
+                - real_pub (t)
+                - real_pub::Pub (t)
+                lib:
+                - Pub (t)
+                - Pub2 (t)
+                - Pub2 (v)
+            "#]],
         );
-
-        assert_snapshot!(map, @r###"
-        main:
-        - publ1 (t)
-        - real_pu2 (t)
-        - real_pub (t)
-        - real_pub::Pub (t)
-        lib:
-        - Pub (t)
-        - Pub2 (t)
-        - Pub2 (v)
-        "###);
     }
 
     #[test]
     fn prefers_shortest_path() {
-        let map = import_map(
+        check(
             r"
             //- /main.rs crate:main
 
@@ -409,21 +466,20 @@ mod tests {
                 pub use super::sub::subsub::Def;
             }
         ",
+            expect![[r#"
+                main:
+                - sub (t)
+                - sub::Def (t)
+                - sub::subsub (t)
+            "#]],
         );
-
-        assert_snapshot!(map, @r###"
-        main:
-        - sub (t)
-        - sub::Def (t)
-        - sub::subsub (t)
-        "###);
     }
 
     #[test]
     fn type_reexport_cross_crate() {
         // Reexports need to be visible from a crate, even if the original crate exports the item
         // at a shorter path.
-        let map = import_map(
+        check(
             r"
             //- /main.rs crate:main deps:lib
             pub mod m {
@@ -432,22 +488,21 @@ mod tests {
             //- /lib.rs crate:lib
             pub struct S;
         ",
+            expect![[r#"
+                main:
+                - m (t)
+                - m::S (t)
+                - m::S (v)
+                lib:
+                - S (t)
+                - S (v)
+            "#]],
         );
-
-        assert_snapshot!(map, @r###"
-        main:
-        - m (t)
-        - m::S (t)
-        - m::S (v)
-        lib:
-        - S (t)
-        - S (v)
-        "###);
     }
 
     #[test]
     fn macro_reexport() {
-        let map = import_map(
+        check(
             r"
             //- /main.rs crate:main deps:lib
             pub mod m {
@@ -459,21 +514,20 @@ mod tests {
                 () => {};
             }
         ",
+            expect![[r#"
+                main:
+                - m (t)
+                - m::pub_macro (m)
+                lib:
+                - pub_macro (m)
+            "#]],
         );
-
-        assert_snapshot!(map, @r###"
-        main:
-        - m (t)
-        - m::pub_macro (m)
-        lib:
-        - pub_macro (m)
-        "###);
     }
 
     #[test]
     fn module_reexport() {
         // Reexporting modules from a dependency adds all contents to the import map.
-        let map = import_map(
+        check(
             r"
             //- /main.rs crate:main deps:lib
             pub use lib::module as reexported_module;
@@ -482,24 +536,23 @@ mod tests {
                 pub struct S;
             }
         ",
+            expect![[r#"
+                main:
+                - reexported_module (t)
+                - reexported_module::S (t)
+                - reexported_module::S (v)
+                lib:
+                - module (t)
+                - module::S (t)
+                - module::S (v)
+            "#]],
         );
-
-        assert_snapshot!(map, @r###"
-        main:
-        - reexported_module (t)
-        - reexported_module::S (t)
-        - reexported_module::S (v)
-        lib:
-        - module (t)
-        - module::S (t)
-        - module::S (v)
-        "###);
     }
 
     #[test]
     fn cyclic_module_reexport() {
         // A cyclic reexport does not hang.
-        let map = import_map(
+        check(
             r"
             //- /lib.rs crate:lib
             pub mod module {
@@ -511,36 +564,35 @@ mod tests {
                 pub use super::module;
             }
         ",
+            expect![[r#"
+                lib:
+                - module (t)
+                - module::S (t)
+                - module::S (v)
+                - sub (t)
+            "#]],
         );
-
-        assert_snapshot!(map, @r###"
-        lib:
-        - module (t)
-        - module::S (t)
-        - module::S (v)
-        - sub (t)
-        "###);
     }
 
     #[test]
     fn private_macro() {
-        let map = import_map(
+        check(
             r"
             //- /lib.rs crate:lib
             macro_rules! private_macro {
                 () => {};
             }
         ",
-        );
+            expect![[r#"
+                lib:
 
-        assert_snapshot!(map, @r###"
-        lib:
-        "###);
+            "#]],
+        );
     }
 
     #[test]
     fn namespacing() {
-        let map = import_map(
+        check(
             r"
             //- /lib.rs crate:lib
             pub struct Thing;     // t + v
@@ -549,16 +601,15 @@ mod tests {
                 () => {};
             }
         ",
+            expect![[r#"
+                lib:
+                - Thing (m)
+                - Thing (t)
+                - Thing (v)
+            "#]],
         );
 
-        assert_snapshot!(map, @r###"
-        lib:
-        - Thing (m)
-        - Thing (t)
-        - Thing (v)
-        "###);
-
-        let map = import_map(
+        check(
             r"
             //- /lib.rs crate:lib
             pub mod Thing {}      // t
@@ -567,13 +618,12 @@ mod tests {
                 () => {};
             }
         ",
+            expect![[r#"
+                lib:
+                - Thing (m)
+                - Thing (t)
+            "#]],
         );
-
-        assert_snapshot!(map, @r###"
-        lib:
-        - Thing (m)
-        - Thing (t)
-        "###);
     }
 
     #[test]
@@ -602,23 +652,33 @@ mod tests {
             }
         "#;
 
-        let res = search_dependencies_of(ra_fixture, "main", Query::new("fmt"));
-        assert_snapshot!(res, @r###"
-        dep::fmt (t)
-        dep::Fmt (t)
-        dep::Fmt (v)
-        dep::Fmt (m)
-        dep::fmt::Display (t)
-        dep::format (v)
-        "###);
-
-        let res = search_dependencies_of(ra_fixture, "main", Query::new("fmt").anchor_end());
-        assert_snapshot!(res, @r###"
-        dep::fmt (t)
-        dep::Fmt (t)
-        dep::Fmt (v)
-        dep::Fmt (m)
-        "###);
+        check_search(
+            ra_fixture,
+            "main",
+            Query::new("fmt"),
+            expect![[r#"
+                dep::fmt (t)
+                dep::Fmt (t)
+                dep::Fmt (v)
+                dep::Fmt (m)
+                dep::fmt::Display (t)
+                dep::format (v)
+                dep::fmt::Display (t)
+            "#]],
+        );
+
+        check_search(
+            ra_fixture,
+            "main",
+            Query::new("fmt").anchor_end(),
+            expect![[r#"
+                dep::fmt (t)
+                dep::Fmt (t)
+                dep::Fmt (v)
+                dep::Fmt (m)
+                dep::fmt::Display (t)
+            "#]],
+        );
     }
 
     #[test]
@@ -631,26 +691,32 @@ mod tests {
             pub struct FMT;
         "#;
 
-        let res = search_dependencies_of(ra_fixture, "main", Query::new("FMT"));
-
-        assert_snapshot!(res, @r###"
-        dep::fmt (t)
-        dep::fmt (v)
-        dep::FMT (t)
-        dep::FMT (v)
-        "###);
-
-        let res = search_dependencies_of(ra_fixture, "main", Query::new("FMT").case_sensitive());
+        check_search(
+            ra_fixture,
+            "main",
+            Query::new("FMT"),
+            expect![[r#"
+                dep::fmt (t)
+                dep::fmt (v)
+                dep::FMT (t)
+                dep::FMT (v)
+            "#]],
+        );
 
-        assert_snapshot!(res, @r###"
-        dep::FMT (t)
-        dep::FMT (v)
-        "###);
+        check_search(
+            ra_fixture,
+            "main",
+            Query::new("FMT").case_sensitive(),
+            expect![[r#"
+                dep::FMT (t)
+                dep::FMT (v)
+            "#]],
+        );
     }
 
     #[test]
     fn search_limit() {
-        let res = search_dependencies_of(
+        check_search(
             r#"
         //- /main.rs crate:main deps:dep
         //- /dep.rs crate:dep
@@ -670,10 +736,10 @@ mod tests {
     "#,
             "main",
             Query::new("").limit(2),
+            expect![[r#"
+                dep::fmt (t)
+                dep::Fmt (t)
+            "#]],
         );
-        assert_snapshot!(res, @r###"
-        dep::fmt (t)
-        dep::Fmt (t)
-        "###);
     }
 }