Merge #4819

4819: Add an FST index to `ImportMap` and use it to speed up auto import r=matklad a=jonas-schievink

For the importing crate, we still use the symbol index, but I've modified it to only look at files that comprise that crate (instead of the whole workspace).

Oh, and since now the symbol query limit is respected correctly, it's possible that some results from the local crate now disappear if there are many matches.

Fixes https://github.com/rust-analyzer/rust-analyzer/issues/4763

Co-authored-by: Jonas Schievink <jonas.schievink@ferrous-systems.com>

Cargo.lock

@@ -981,7 +981,10 @@ dependencies = [
  "anymap",
  "drop_bomb",
  "either",
+ "fst",
+ "indexmap",
  "insta",
+ "itertools",
  "log",
  "once_cell",
  "ra_arena",

crates/ra_assists/src/handlers/auto_import.rs

@@ -130,7 +130,7 @@ impl AutoImportAssets {
     fn search_for_imports(&self, db: &RootDatabase) -> BTreeSet<ModPath> {
         let _p = profile("auto_import::search_for_imports");
         let current_crate = self.module_with_name_to_import.krate();
-        ImportsLocator::new(db)
+        ImportsLocator::new(db, current_crate)
             .find_imports(&self.get_search_query())
             .into_iter()
             .filter_map(|candidate| match &self.import_candidate {
@@ -841,4 +841,105 @@ fn main() {
             ",
         )
     }
+
+    #[test]
+    fn dep_import() {
+        check_assist(
+            auto_import,
+            r"
+                    //- /lib.rs crate:dep
+                    pub struct Struct;
+
+                    //- /main.rs crate:main deps:dep
+                    fn main() {
+                        Struct<|>
+                    }",
+            r"use dep::Struct;
+
+fn main() {
+    Struct
+}
+",
+        );
+    }
+
+    #[test]
+    fn whole_segment() {
+        // Tests that only imports whose last segment matches the identifier get suggested.
+        check_assist(
+            auto_import,
+            r"
+                    //- /lib.rs crate:dep
+                    pub mod fmt {
+                        pub trait Display {}
+                    }
+
+                    pub fn panic_fmt() {}
+
+                    //- /main.rs crate:main deps:dep
+                    struct S;
+
+                    impl f<|>mt::Display for S {}",
+            r"use dep::fmt;
+
+struct S;
+impl fmt::Display for S {}
+",
+        );
+    }
+
+    #[test]
+    fn macro_generated() {
+        // Tests that macro-generated items are suggested from external crates.
+        check_assist(
+            auto_import,
+            r"
+                    //- /lib.rs crate:dep
+
+                    macro_rules! mac {
+                        () => {
+                            pub struct Cheese;
+                        };
+                    }
+
+                    mac!();
+
+                    //- /main.rs crate:main deps:dep
+
+                    fn main() {
+                        Cheese<|>;
+                    }",
+            r"use dep::Cheese;
+
+fn main() {
+    Cheese;
+}
+",
+        );
+    }
+
+    #[test]
+    fn casing() {
+        // Tests that differently cased names don't interfere and we only suggest the matching one.
+        check_assist(
+            auto_import,
+            r"
+                    //- /lib.rs crate:dep
+
+                    pub struct FMT;
+                    pub struct fmt;
+
+                    //- /main.rs crate:main deps:dep
+
+                    fn main() {
+                        FMT<|>;
+                    }",
+            r"use dep::FMT;
+
+fn main() {
+    FMT;
+}
+",
+        );
+    }
 }

crates/ra_hir/src/code_model.rs

@@ -9,6 +9,7 @@ use hir_def::{
     builtin_type::BuiltinType,
     docs::Documentation,
     expr::{BindingAnnotation, Pat, PatId},
+    import_map,
     per_ns::PerNs,
     resolver::{HasResolver, Resolver},
     type_ref::{Mutability, TypeRef},
@@ -98,6 +99,23 @@ impl Crate {
         db.crate_graph()[self.id].display_name.as_ref().cloned()
     }
 
+    pub fn query_external_importables(
+        self,
+        db: &dyn DefDatabase,
+        query: &str,
+    ) -> impl Iterator<Item = Either<ModuleDef, MacroDef>> {
+        import_map::search_dependencies(
+            db,
+            self.into(),
+            import_map::Query::new(query).anchor_end().case_sensitive().limit(40),
+        )
+        .into_iter()
+        .map(|item| match item {
+            ItemInNs::Types(mod_id) | ItemInNs::Values(mod_id) => Either::Left(mod_id.into()),
+            ItemInNs::Macros(mac_id) => Either::Right(mac_id.into()),
+        })
+    }
+
     pub fn all(db: &dyn HirDatabase) -> Vec<Crate> {
         db.crate_graph().iter().map(|id| Crate { id }).collect()
     }

crates/ra_hir_def/Cargo.toml

@@ -14,6 +14,9 @@ rustc-hash = "1.1.0"
 either = "1.5.3"
 anymap = "0.12.1"
 drop_bomb = "0.1.4"
+fst = { version = "0.4", default-features = false }
+itertools = "0.9.0"
+indexmap = "1.4.0"
 
 stdx = { path = "../stdx" }
 

crates/ra_hir_def/src/import_map.rs

@@ -1,9 +1,11 @@
 //! A map of all publicly exported items in a crate.
 
-use std::{collections::hash_map::Entry, fmt, sync::Arc};
+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::FxHashMap;
+use rustc_hash::FxHasher;
 
 use crate::{
     db::DefDatabase,
@@ -13,6 +15,8 @@ use crate::{
     ModuleDefId, ModuleId,
 };
 
+type FxIndexMap<K, V> = IndexMap<K, V, BuildHasherDefault<FxHasher>>;
+
 /// A map from publicly exported items to the path needed to import/name them from a downstream
 /// crate.
 ///
@@ -21,16 +25,24 @@ use crate::{
 ///
 /// 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(Eq, PartialEq)]
 pub struct ImportMap {
-    map: FxHashMap<ItemInNs, ModPath>,
+    map: FxIndexMap<ItemInNs, ModPath>,
+
+    /// List of keys stored in `map`, sorted lexicographically by their `ModPath`. Indexed by the
+    /// values returned by running `fst`.
+    ///
+    /// Since a path can refer to multiple items due to namespacing, we store all items with the
+    /// same path right after each other. This allows us to find all items after the FST gives us
+    /// the index of the first one.
+    importables: Vec<ItemInNs>,
+    fst: fst::Map<Vec<u8>>,
 }
 
 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 = FxHashMap::with_capacity_and_hasher(64, Default::default());
+        let mut import_map = FxIndexMap::with_capacity_and_hasher(64, Default::default());
 
         // We look only into modules that are public(ly reexported), starting with the crate root.
         let empty = ModPath { kind: PathKind::Plain, segments: vec![] };
@@ -88,7 +100,34 @@ impl ImportMap {
             }
         }
 
-        Arc::new(Self { map: import_map })
+        let mut importables = import_map.iter().collect::<Vec<_>>();
+
+        importables.sort_by(cmp);
+
+        // Build the FST, taking care not to insert duplicate values.
+
+        let mut builder = fst::MapBuilder::memory();
+        let mut last_batch_start = 0;
+
+        for idx in 0..importables.len() {
+            if let Some(next_item) = importables.get(idx + 1) {
+                if cmp(&importables[last_batch_start], next_item) == Ordering::Equal {
+                    continue;
+                }
+            }
+
+            let start = last_batch_start;
+            last_batch_start = idx + 1;
+
+            let key = fst_path(&importables[start].1);
+
+            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();
+
+        Arc::new(Self { map: import_map, fst, importables })
     }
 
     /// Returns the `ModPath` needed to import/mention `item`, relative to this crate's root.
@@ -97,6 +136,15 @@ impl ImportMap {
     }
 }
 
+impl PartialEq for ImportMap {
+    fn eq(&self, other: &Self) -> bool {
+        // `fst` and `importables` are built from `map`, so we don't need to compare them.
+        self.map == other.map
+    }
+}
+
+impl Eq for ImportMap {}
+
 impl fmt::Debug for ImportMap {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut importable_paths: Vec<_> = self
@@ -117,19 +165,135 @@ impl fmt::Debug for ImportMap {
     }
 }
 
+fn fst_path(path: &ModPath) -> String {
+    let mut s = path.to_string();
+    s.make_ascii_lowercase();
+    s
+}
+
+fn cmp((_, lhs): &(&ItemInNs, &ModPath), (_, rhs): &(&ItemInNs, &ModPath)) -> Ordering {
+    let lhs_str = fst_path(lhs);
+    let rhs_str = fst_path(rhs);
+    lhs_str.cmp(&rhs_str)
+}
+
+#[derive(Debug)]
+pub struct Query {
+    query: String,
+    lowercased: String,
+    anchor_end: bool,
+    case_sensitive: bool,
+    limit: usize,
+}
+
+impl Query {
+    pub fn new(query: &str) -> Self {
+        Self {
+            lowercased: query.to_lowercase(),
+            query: query.to_string(),
+            anchor_end: false,
+            case_sensitive: false,
+            limit: usize::max_value(),
+        }
+    }
+
+    /// Only returns items whose paths end with the (case-insensitive) query string as their last
+    /// segment.
+    pub fn anchor_end(self) -> Self {
+        Self { anchor_end: true, ..self }
+    }
+
+    /// Limits the returned number of items to `limit`.
+    pub fn limit(self, limit: usize) -> Self {
+        Self { limit, ..self }
+    }
+
+    /// Respect casing of the query string when matching.
+    pub fn case_sensitive(self) -> Self {
+        Self { case_sensitive: true, ..self }
+    }
+}
+
+/// Searches dependencies of `krate` for an importable path matching `query`.
+///
+/// This returns a list of items that could be imported from dependencies of `krate`.
+pub fn search_dependencies<'a>(
+    db: &'a dyn DefDatabase,
+    krate: CrateId,
+    query: Query,
+) -> Vec<ItemInNs> {
+    let _p = ra_prof::profile("search_dependencies").detail(|| format!("{:?}", query));
+
+    let graph = db.crate_graph();
+    let import_maps: Vec<_> =
+        graph[krate].dependencies.iter().map(|dep| db.import_map(dep.crate_id)).collect();
+
+    let automaton = fst::automaton::Subsequence::new(&query.lowercased);
+
+    let mut op = fst::map::OpBuilder::new();
+    for map in &import_maps {
+        op = op.add(map.fst.search(&automaton));
+    }
+
+    let mut stream = op.union();
+    let mut res = Vec::new();
+    while let Some((_, indexed_values)) = stream.next() {
+        for indexed_value in indexed_values {
+            let import_map = &import_maps[indexed_value.index];
+            let importables = &import_map.importables[indexed_value.value as usize..];
+
+            // Path shared by the importable items in this group.
+            let path = &import_map.map[&importables[0]];
+
+            if query.anchor_end {
+                // Last segment must match query.
+                let last = path.segments.last().unwrap().to_string();
+                if last.to_lowercase() != query.lowercased {
+                    continue;
+                }
+            }
+
+            // Add the items from this `ModPath` group. Those are all subsequent items in
+            // `importables` whose paths match `path`.
+            let iter = importables.iter().copied().take_while(|item| {
+                let item_path = &import_map.map[item];
+                fst_path(item_path) == fst_path(path)
+            });
+
+            if query.case_sensitive {
+                // FIXME: This does not do a subsequence match.
+                res.extend(iter.filter(|item| {
+                    let item_path = &import_map.map[item];
+                    item_path.to_string().contains(&query.query)
+                }));
+            } else {
+                res.extend(iter);
+            }
+
+            if res.len() >= query.limit {
+                res.truncate(query.limit);
+                return res;
+            }
+        }
+    }
+
+    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;
+    use ra_db::{SourceDatabase, Upcast};
 
     fn import_map(ra_fixture: &str) -> String {
         let db = TestDB::with_files(ra_fixture);
         let crate_graph = db.crate_graph();
 
-        let import_maps: Vec<_> = crate_graph
+        let s = crate_graph
             .iter()
             .filter_map(|krate| {
                 let cdata = &crate_graph[krate];
@@ -139,9 +303,41 @@ mod tests {
 
                 Some(format!("{}:\n{:?}", name, map))
             })
-            .collect();
+            .join("\n");
+        s
+    }
 
-        import_maps.join("\n")
+    fn search_dependencies_of(ra_fixture: &str, krate_name: &str, query: Query) -> String {
+        let db = TestDB::with_files(ra_fixture);
+        let crate_graph = db.crate_graph();
+        let krate = crate_graph
+            .iter()
+            .find(|krate| {
+                crate_graph[*krate].display_name.as_ref().map(|n| n.to_string())
+                    == Some(krate_name.to_string())
+            })
+            .unwrap();
+
+        search_dependencies(db.upcast(), krate, query)
+            .into_iter()
+            .filter_map(|item| {
+                let mark = match item {
+                    ItemInNs::Types(_) => "t",
+                    ItemInNs::Values(_) => "v",
+                    ItemInNs::Macros(_) => "m",
+                };
+                item.krate(db.upcast()).map(|krate| {
+                    let map = db.import_map(krate);
+                    let path = map.path_of(item).unwrap();
+                    format!(
+                        "{}::{} ({})",
+                        crate_graph[krate].display_name.as_ref().unwrap(),
+                        path,
+                        mark
+                    )
+                })
+            })
+            .join("\n")
     }
 
     #[test]
@@ -328,4 +524,143 @@ mod tests {
         lib:
         "###);
     }
+
+    #[test]
+    fn namespacing() {
+        let map = import_map(
+            r"
+            //- /lib.rs crate:lib
+            pub struct Thing;     // t + v
+            #[macro_export]
+            macro_rules! Thing {  // m
+                () => {};
+            }
+        ",
+        );
+
+        assert_snapshot!(map, @r###"
+        lib:
+        - Thing (m)
+        - Thing (t)
+        - Thing (v)
+        "###);
+
+        let map = import_map(
+            r"
+            //- /lib.rs crate:lib
+            pub mod Thing {}      // t
+            #[macro_export]
+            macro_rules! Thing {  // m
+                () => {};
+            }
+        ",
+        );
+
+        assert_snapshot!(map, @r###"
+        lib:
+        - Thing (m)
+        - Thing (t)
+        "###);
+    }
+
+    #[test]
+    fn search() {
+        let ra_fixture = r#"
+            //- /main.rs crate:main deps:dep
+            //- /dep.rs crate:dep deps:tdep
+            use tdep::fmt as fmt_dep;
+            pub mod fmt {
+                pub trait Display {
+                    fn fmt();
+                }
+            }
+            #[macro_export]
+            macro_rules! Fmt {
+                () => {};
+            }
+            pub struct Fmt;
+
+            pub fn format() {}
+            pub fn no() {}
+
+            //- /tdep.rs crate:tdep
+            pub mod fmt {
+                pub struct NotImportableFromMain;
+            }
+        "#;
+
+        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)
+        "###);
+    }
+
+    #[test]
+    fn search_casing() {
+        let ra_fixture = r#"
+            //- /main.rs crate:main deps:dep
+            //- /dep.rs crate:dep
+
+            pub struct fmt;
+            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());
+
+        assert_snapshot!(res, @r###"
+        dep::FMT (t)
+        dep::FMT (v)
+        "###);
+    }
+
+    #[test]
+    fn search_limit() {
+        let res = search_dependencies_of(
+            r#"
+        //- /main.rs crate:main deps:dep
+        //- /dep.rs crate:dep
+        pub mod fmt {
+            pub trait Display {
+                fn fmt();
+            }
+        }
+        #[macro_export]
+        macro_rules! Fmt {
+            () => {};
+        }
+        pub struct Fmt;
+
+        pub fn format() {}
+        pub fn no() {}
+    "#,
+            "main",
+            Query::new("").limit(2),
+        );
+        assert_snapshot!(res, @r###"
+        dep::fmt (t)
+        dep::Fmt (t)
+        "###);
+    }
 }

crates/ra_ide_db/src/imports_locator.rs

@@ -1,7 +1,7 @@
 //! This module contains an import search funcionality that is provided to the ra_assists module.
 //! Later, this should be moved away to a separate crate that is accessible from the ra_assists module.
 
-use hir::{MacroDef, ModuleDef, Semantics};
+use hir::{Crate, MacroDef, ModuleDef, Semantics};
 use ra_prof::profile;
 use ra_syntax::{ast, AstNode, SyntaxKind::NAME};
 
@@ -11,44 +11,46 @@ use crate::{
     RootDatabase,
 };
 use either::Either;
+use rustc_hash::FxHashSet;
 
 pub struct ImportsLocator<'a> {
     sema: Semantics<'a, RootDatabase>,
+    krate: Crate,
 }
 
 impl<'a> ImportsLocator<'a> {
-    pub fn new(db: &'a RootDatabase) -> Self {
-        Self { sema: Semantics::new(db) }
+    pub fn new(db: &'a RootDatabase, krate: Crate) -> Self {
+        Self { sema: Semantics::new(db), krate }
     }
 
     pub fn find_imports(&mut self, name_to_import: &str) -> Vec<Either<ModuleDef, MacroDef>> {
         let _p = profile("search_for_imports");
         let db = self.sema.db;
 
-        let project_results = {
+        // Query dependencies first.
+        let mut candidates: FxHashSet<_> =
+            self.krate.query_external_importables(db, name_to_import).collect();
+
+        // Query the local crate using the symbol index.
+        let local_results = {
             let mut query = Query::new(name_to_import.to_string());
             query.exact();
             query.limit(40);
-            symbol_index::world_symbols(db, query)
-        };
-        let lib_results = {
-            let mut query = Query::new(name_to_import.to_string());
-            query.libs();
-            query.exact();
-            query.limit(40);
-            symbol_index::world_symbols(db, query)
+            symbol_index::crate_symbols(db, self.krate.into(), query)
         };
 
-        project_results
-            .into_iter()
-            .chain(lib_results.into_iter())
-            .filter_map(|import_candidate| self.get_name_definition(&import_candidate))
-            .filter_map(|name_definition_to_import| match name_definition_to_import {
-                Definition::ModuleDef(module_def) => Some(Either::Left(module_def)),
-                Definition::Macro(macro_def) => Some(Either::Right(macro_def)),
-                _ => None,
-            })
-            .collect()
+        candidates.extend(
+            local_results
+                .into_iter()
+                .filter_map(|import_candidate| self.get_name_definition(&import_candidate))
+                .filter_map(|name_definition_to_import| match name_definition_to_import {
+                    Definition::ModuleDef(module_def) => Some(Either::Left(module_def)),
+                    Definition::Macro(macro_def) => Some(Either::Right(macro_def)),
+                    _ => None,
+                }),
+        );
+
+        candidates.into_iter().collect()
     }
 
     fn get_name_definition(&mut self, import_candidate: &FileSymbol) -> Option<Definition> {

crates/ra_ide_db/src/symbol_index.rs

@@ -29,9 +29,10 @@ use std::{
 };
 
 use fst::{self, Streamer};
+use hir::db::DefDatabase;
 use ra_db::{
     salsa::{self, ParallelDatabase},
-    FileId, SourceDatabaseExt, SourceRootId,
+    CrateId, FileId, SourceDatabaseExt, SourceRootId,
 };
 use ra_syntax::{
     ast::{self, NameOwner},
@@ -110,6 +111,14 @@ fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex>
     Arc::new(SymbolIndex::new(symbols))
 }
 
+/// Need to wrap Snapshot to provide `Clone` impl for `map_with`
+struct Snap(salsa::Snapshot<RootDatabase>);
+impl Clone for Snap {
+    fn clone(&self) -> Snap {
+        Snap(self.0.snapshot())
+    }
+}
+
 // Feature: Workspace Symbol
 //
 // Uses fuzzy-search to find types, modules and functions by name across your
@@ -132,13 +141,7 @@ fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex>
 // | VS Code | kbd:[Ctrl+T]
 // |===
 pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
-    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
-    struct Snap(salsa::Snapshot<RootDatabase>);
-    impl Clone for Snap {
-        fn clone(&self) -> Snap {
-            Snap(self.0.snapshot())
-        }
-    }
+    let _p = ra_prof::profile("world_symbols").detail(|| query.query.clone());
 
     let buf: Vec<Arc<SymbolIndex>> = if query.libs {
         let snap = Snap(db.snapshot());
@@ -173,6 +176,33 @@ pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
     query.search(&buf)
 }
 
+pub fn crate_symbols(db: &RootDatabase, krate: CrateId, query: Query) -> Vec<FileSymbol> {
+    // FIXME(#4842): This now depends on CrateDefMap, why not build the entire symbol index from
+    // that instead?
+
+    let def_map = db.crate_def_map(krate);
+    let mut files = Vec::new();
+    let mut modules = vec![def_map.root];
+    while let Some(module) = modules.pop() {
+        let data = &def_map[module];
+        files.extend(data.origin.file_id());
+        modules.extend(data.children.values());
+    }
+
+    let snap = Snap(db.snapshot());
+
+    #[cfg(not(feature = "wasm"))]
+    let buf = files
+        .par_iter()
+        .map_with(snap, |db, &file_id| db.0.file_symbols(file_id))
+        .collect::<Vec<_>>();
+
+    #[cfg(feature = "wasm")]
+    let buf = files.iter().map(|&file_id| snap.0.file_symbols(file_id)).collect::<Vec<_>>();
+
+    query.search(&buf)
+}
+
 pub fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
     let name = name_ref.text();
     let mut query = Query::new(name.to_string());
@@ -298,9 +328,6 @@ impl Query {
         let mut stream = op.union();
         let mut res = Vec::new();
         while let Some((_, indexed_values)) = stream.next() {
-            if res.len() >= self.limit {
-                break;
-            }
             for indexed_value in indexed_values {
                 let symbol_index = &indices[indexed_value.index];
                 let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);
@@ -312,7 +339,11 @@ impl Query {
                     if self.exact && symbol.name != self.query {
                         continue;
                     }
+
                     res.push(symbol.clone());
+                    if res.len() >= self.limit {
+                        return res;
+                    }
                 }
             }
         }