Restructure symbol queries

crates/hir/src/symbols.rs

@@ -102,21 +102,33 @@ pub struct SymbolCollector<'a> {
 /// Given a [`ModuleId`] and a [`HirDatabase`], use the DefMap for the module's crate to collect
 /// all symbols that should be indexed for the given module.
 impl<'a> SymbolCollector<'a> {
-    pub fn collect(db: &dyn HirDatabase, module: Module) -> Vec<FileSymbol> {
-        let mut symbol_collector = SymbolCollector {
+    pub fn new(db: &'a dyn HirDatabase) -> Self {
+        SymbolCollector {
             db,
             symbols: Default::default(),
+            work: Default::default(),
             current_container_name: None,
-            // The initial work is the root module we're collecting, additional work will
-            // be populated as we traverse the module's definitions.
-            work: vec![SymbolCollectorWork { module_id: module.into(), parent: None }],
-        };
-
-        while let Some(work) = symbol_collector.work.pop() {
-            symbol_collector.do_work(work);
         }
+    }
 
-        symbol_collector.symbols
+    pub fn collect(&mut self, module: Module) {
+        // The initial work is the root module we're collecting, additional work will
+        // be populated as we traverse the module's definitions.
+        self.work.push(SymbolCollectorWork { module_id: module.into(), parent: None });
+
+        while let Some(work) = self.work.pop() {
+            self.do_work(work);
+        }
+    }
+
+    pub fn finish(self) -> Vec<FileSymbol> {
+        self.symbols
+    }
+
+    pub fn collect_module(db: &dyn HirDatabase, module: Module) -> Vec<FileSymbol> {
+        let mut symbol_collector = SymbolCollector::new(db);
+        symbol_collector.collect(module);
+        symbol_collector.finish()
     }
 
     fn do_work(&mut self, work: SymbolCollectorWork) {

crates/ide-db/src/items_locator.rs

@@ -115,7 +115,8 @@ fn find_items<'a>(
         });
 
     // Query the local crate using the symbol index.
-    let local_results = symbol_index::crate_symbols(db, krate, local_query)
+    let local_results = local_query
+        .search(&symbol_index::crate_symbols(db, krate))
         .into_iter()
         .filter_map(move |local_candidate| get_name_definition(sema, &local_candidate))
         .filter_map(|name_definition_to_import| match name_definition_to_import {

crates/ide-db/src/symbol_index.rs

@@ -93,12 +93,15 @@ impl Query {
 pub trait SymbolsDatabase: HirDatabase + SourceDatabaseExt + Upcast<dyn HirDatabase> {
     /// The symbol index for a given module. These modules should only be in source roots that
     /// are inside local_roots.
-    // FIXME: We should probably LRU  this
     fn module_symbols(&self, module: Module) -> Arc<SymbolIndex>;
 
     /// The symbol index for a given source root within library_roots.
     fn library_symbols(&self, source_root_id: SourceRootId) -> Arc<SymbolIndex>;
 
+    #[salsa::transparent]
+    /// The symbol indices of modules that make up a given crate.
+    fn crate_symbols(&self, krate: Crate) -> Box<[Arc<SymbolIndex>]>;
+
     /// The set of "local" (that is, from the current workspace) roots.
     /// Files in local roots are assumed to change frequently.
     #[salsa::input]
@@ -113,26 +116,33 @@ pub trait SymbolsDatabase: HirDatabase + SourceDatabaseExt + Upcast<dyn HirDatab
 fn library_symbols(db: &dyn SymbolsDatabase, source_root_id: SourceRootId) -> Arc<SymbolIndex> {
     let _p = profile::span("library_symbols");
 
-    // todo: this could be parallelized, once I figure out how to do that...
-    let symbols = db
-        .source_root_crates(source_root_id)
+    let mut symbol_collector = SymbolCollector::new(db.upcast());
+
+    db.source_root_crates(source_root_id)
         .iter()
         .flat_map(|&krate| Crate::from(krate).modules(db.upcast()))
-        // we specifically avoid calling SymbolsDatabase::module_symbols here, even they do the same thing,
+        // we specifically avoid calling other SymbolsDatabase queries here, even though they do the same thing,
         // as the index for a library is not going to really ever change, and we do not want to store each
-        // module's index in salsa.
-        .flat_map(|module| SymbolCollector::collect(db.upcast(), module))
-        .collect();
+        // the module or crate indices for those in salsa unless we need to.
+        .for_each(|module| symbol_collector.collect(module));
 
+    let mut symbols = symbol_collector.finish();
+    symbols.shrink_to_fit();
     Arc::new(SymbolIndex::new(symbols))
 }
 
 fn module_symbols(db: &dyn SymbolsDatabase, module: Module) -> Arc<SymbolIndex> {
     let _p = profile::span("module_symbols");
-    let symbols = SymbolCollector::collect(db.upcast(), module);
+
+    let symbols = SymbolCollector::collect_module(db.upcast(), module);
     Arc::new(SymbolIndex::new(symbols))
 }
 
+pub fn crate_symbols(db: &dyn SymbolsDatabase, krate: Crate) -> Box<[Arc<SymbolIndex>]> {
+    let _p = profile::span("crate_symbols");
+    krate.modules(db.upcast()).into_iter().map(|module| db.module_symbols(module)).collect()
+}
+
 /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
 struct Snap<DB>(DB);
 impl<DB: ParallelDatabase> Snap<salsa::Snapshot<DB>> {
@@ -188,36 +198,21 @@ pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
             .map_with(Snap::new(db), |snap, &root| snap.library_symbols(root))
             .collect()
     } else {
-        let mut modules = Vec::new();
+        let mut crates = Vec::new();
 
         for &root in db.local_roots().iter() {
-            let crates = db.source_root_crates(root);
-            for &krate in crates.iter() {
-                modules.extend(Crate::from(krate).modules(db));
-            }
+            crates.extend(db.source_root_crates(root).iter().copied())
         }
-
-        modules
-            .par_iter()
-            .map_with(Snap::new(db), |snap, &module| snap.module_symbols(module))
-            .collect()
+        let indices: Vec<_> = crates
+            .into_par_iter()
+            .map_with(Snap::new(db), |snap, krate| snap.crate_symbols(krate.into()))
+            .collect();
+        indices.iter().flat_map(|indices| indices.iter().cloned()).collect()
     };
 
     query.search(&indices)
 }
 
-pub fn crate_symbols(db: &RootDatabase, krate: Crate, query: Query) -> Vec<FileSymbol> {
-    let _p = profile::span("crate_symbols").detail(|| format!("{query:?}"));
-
-    let modules = krate.modules(db);
-    let indices: Vec<_> = modules
-        .par_iter()
-        .map_with(Snap::new(db), |snap, &module| snap.module_symbols(module))
-        .collect();
-
-    query.search(&indices)
-}
-
 #[derive(Default)]
 pub struct SymbolIndex {
     symbols: Vec<FileSymbol>,
@@ -275,7 +270,12 @@ impl SymbolIndex {
             builder.insert(key, value).unwrap();
         }
 
-        let map = fst::Map::new(builder.into_inner().unwrap()).unwrap();
+        let map = fst::Map::new({
+            let mut buf = builder.into_inner().unwrap();
+            buf.shrink_to_fit();
+            buf
+        })
+        .unwrap();
         SymbolIndex { symbols, map }
     }
 
@@ -419,7 +419,7 @@ struct StructInModB;
             .modules(&db)
             .into_iter()
             .map(|module_id| {
-                let mut symbols = SymbolCollector::collect(&db, module_id);
+                let mut symbols = SymbolCollector::collect_module(&db, module_id);
                 symbols.sort_by_key(|it| it.name.clone());
                 (module_id, symbols)
             })