Record and suggest trait items via ImportMap

2 files changed, 179 insertions, 8 deletions

diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index 7b6499a08..4cd77adbf 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -811,6 +811,146 @@ fn main() {
     }
 
     #[test]
+    fn trait_method_cross_crate() {
+        check_assist(
+            auto_import,
+            r"
+            //- /main.rs crate:main deps:dep
+            fn main() {
+                let test_struct = dep::test_mod::TestStruct {};
+                test_struct.test_meth<|>od()
+            }
+            //- /dep.rs crate:dep
+            pub mod test_mod {
+                pub trait TestTrait {
+                    fn test_method(&self);
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_method(&self) {}
+                }
+            }
+            ",
+            r"
+            use dep::test_mod::TestTrait;
+
+            fn main() {
+                let test_struct = dep::test_mod::TestStruct {};
+                test_struct.test_method()
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn assoc_fn_cross_crate() {
+        check_assist(
+            auto_import,
+            r"
+            //- /main.rs crate:main deps:dep
+            fn main() {
+                dep::test_mod::TestStruct::test_func<|>tion
+            }
+            //- /dep.rs crate:dep
+            pub mod test_mod {
+                pub trait TestTrait {
+                    fn test_function();
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_function() {}
+                }
+            }
+            ",
+            r"
+            use dep::test_mod::TestTrait;
+
+            fn main() {
+                dep::test_mod::TestStruct::test_function
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn assoc_const_cross_crate() {
+        check_assist(
+            auto_import,
+            r"
+            //- /main.rs crate:main deps:dep
+            fn main() {
+                dep::test_mod::TestStruct::CONST<|>
+            }
+            //- /dep.rs crate:dep
+            pub mod test_mod {
+                pub trait TestTrait {
+                    const CONST: bool;
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    const CONST: bool = true;
+                }
+            }
+            ",
+            r"
+            use dep::test_mod::TestTrait;
+
+            fn main() {
+                dep::test_mod::TestStruct::CONST
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn assoc_fn_as_method_cross_crate() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+            //- /main.rs crate:main deps:dep
+            fn main() {
+                let test_struct = dep::test_mod::TestStruct {};
+                test_struct.test_func<|>tion()
+            }
+            //- /dep.rs crate:dep
+            pub mod test_mod {
+                pub trait TestTrait {
+                    fn test_function();
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_function() {}
+                }
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn private_trait_cross_crate() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+            //- /main.rs crate:main deps:dep
+            fn main() {
+                let test_struct = dep::test_mod::TestStruct {};
+                test_struct.test_meth<|>od()
+            }
+            //- /dep.rs crate:dep
+            pub mod test_mod {
+                trait TestTrait {
+                    fn test_method(&self);
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_method(&self) {}
+                }
+            }
+            ",
+        );
+    }
+
+    #[test]
     fn not_applicable_for_imported_trait_for_method() {
         check_assist_not_applicable(
             auto_import,
diff --git a/crates/ra_hir_def/src/import_map.rs b/crates/ra_hir_def/src/import_map.rs
index 68e20d06b..299fe82a8 100644
--- a/crates/ra_hir_def/src/import_map.rs
+++ b/crates/ra_hir_def/src/import_map.rs
@@ -5,14 +5,15 @@ 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 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 +35,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 +47,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<String, 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 +91,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 +111,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 +144,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 +158,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()).or_default().push(*item);
+        }
+    }
 }
 
 impl PartialEq for ImportMap {
@@ -290,6 +308,19 @@ 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
 }