Auto merge of #16261 - Veykril:import-map-fix, r=Veykril

internal: Move query limits to the caller

Prior we calculated up to `limit` entries from a query, then filtered from that leaving us with less entries than the limit in some cases (which might give odd completion behavior due to items disappearing). This changes it so we filter before checking the limit.

crates/hir-def/src/import_map.rs

@@ -295,7 +295,6 @@ pub struct Query {
     search_mode: SearchMode,
     assoc_mode: AssocSearchMode,
     case_sensitive: bool,
-    limit: usize,
 }
 
 impl Query {
@@ -307,7 +306,6 @@ impl Query {
             search_mode: SearchMode::Exact,
             assoc_mode: AssocSearchMode::Include,
             case_sensitive: false,
-            limit: usize::MAX,
         }
     }
 
@@ -329,11 +327,6 @@ impl Query {
         Self { assoc_mode, ..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 }
@@ -413,6 +406,7 @@ fn search_maps(
                 })
                 // we put all entries with the same lowercased name in a row, so stop once we find a
                 // different name in the importables
+                // FIXME: Consider putting a range into the value: u64 as (u32, u32)?
                 .take_while(|&(_, info, _)| {
                     info.name.to_smol_str().as_bytes().eq_ignore_ascii_case(&key)
                 })
@@ -424,13 +418,32 @@ fn search_maps(
                         return true;
                     }
                     let name = info.name.to_smol_str();
+                    // FIXME: Deduplicate this from ide-db
                     match query.search_mode {
-                        SearchMode::Exact => name == query.query,
+                        SearchMode::Exact => !query.case_sensitive || name == query.query,
-                        SearchMode::Prefix => name.starts_with(&query.query),
+                        SearchMode::Prefix => {
+                            query.query.len() <= name.len() && {
+                                let prefix = &name[..query.query.len() as usize];
+                                if query.case_sensitive {
+                                    prefix == query.query
+                                } else {
+                                    prefix.eq_ignore_ascii_case(&query.query)
+                                }
+                            }
+                        }
                         SearchMode::Fuzzy => {
                             let mut name = &*name;
                             query.query.chars().all(|query_char| {
-                                match name.match_indices(query_char).next() {
+                                let m = if query.case_sensitive {
+                                    name.match_indices(query_char).next()
+                                } else {
+                                    name.match_indices([
+                                        query_char,
+                                        query_char.to_ascii_uppercase(),
+                                    ])
+                                    .next()
+                                };
+                                match m {
                                     Some((index, _)) => {
                                         name = &name[index + 1..];
                                         true
@@ -442,10 +455,6 @@ fn search_maps(
                     }
                 });
             res.extend(iter.map(TupleExt::head));
-
-            if res.len() >= query.limit {
-                return res;
-            }
         }
     }
 
@@ -1015,32 +1024,4 @@ pub mod fmt {
             "#]],
         );
     }
-
-    #[test]
-    fn search_limit() {
-        check_search(
-            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("".to_string()).fuzzy().limit(1),
-            expect![[r#"
-                dep::fmt::Display (t)
-            "#]],
-        );
-    }
 }

crates/hir/src/symbols.rs

@@ -163,7 +163,7 @@ impl<'a> SymbolCollector<'a> {
         }
 
         // Record renamed imports.
-        // In case it imports multiple items under different namespaces we just pick one arbitrarily
+        // FIXME: In case it imports multiple items under different namespaces we just pick one arbitrarily
         // for now.
         for id in scope.imports() {
             let loc = id.import.lookup(self.db.upcast());

crates/ide-assists/src/handlers/replace_derive_with_manual_impl.rs

@@ -74,7 +74,6 @@ pub(crate) fn replace_derive_with_manual_impl(
         current_crate,
         NameToImport::exact_case_sensitive(path.segments().last()?.to_string()),
         items_locator::AssocSearchMode::Exclude,
-        Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT.inner()),
     )
     .filter_map(|item| match item.as_module_def()? {
         ModuleDef::Trait(trait_) => Some(trait_),

crates/ide-completion/src/lib.rs

@@ -256,7 +256,6 @@ pub fn resolve_completion_edits(
             current_crate,
             NameToImport::exact_case_sensitive(imported_name),
             items_locator::AssocSearchMode::Include,
-            Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT.inner()),
         );
         let import = items_with_name
             .filter_map(|candidate| {

crates/ide-db/src/imports/import_assets.rs

@@ -333,12 +333,12 @@ fn path_applicable_imports(
                 //
                 // see also an ignored test under FIXME comment in the qualify_path.rs module
                 AssocSearchMode::Exclude,
-                Some(DEFAULT_QUERY_SEARCH_LIMIT.inner()),
             )
             .filter_map(|item| {
                 let mod_path = mod_path(item)?;
                 Some(LocatedImport::new(mod_path, item, item))
             })
+            .take(DEFAULT_QUERY_SEARCH_LIMIT.inner())
             .collect()
         }
         Some(qualifier) => items_locator::items_with_name(
@@ -346,9 +346,9 @@ fn path_applicable_imports(
             current_crate,
             path_candidate.name.clone(),
             AssocSearchMode::Include,
-            Some(DEFAULT_QUERY_SEARCH_LIMIT.inner()),
         )
         .filter_map(|item| import_for_item(sema.db, mod_path, &qualifier, item))
+        .take(DEFAULT_QUERY_SEARCH_LIMIT.inner())
         .collect(),
     }
 }
@@ -505,7 +505,6 @@ fn trait_applicable_items(
         current_crate,
         trait_candidate.assoc_item_name.clone(),
         AssocSearchMode::AssocItemsOnly,
-        Some(DEFAULT_QUERY_SEARCH_LIMIT.inner()),
     )
     .filter_map(|input| item_as_assoc(db, input))
     .filter_map(|assoc| {
@@ -517,6 +516,7 @@ fn trait_applicable_items(
             Some(assoc_item_trait.into())
         }
     })
+    .take(DEFAULT_QUERY_SEARCH_LIMIT.inner())
     .collect();
 
     let mut located_imports = FxHashSet::default();

crates/ide-db/src/items_locator.rs

@@ -19,26 +19,24 @@ pub fn items_with_name<'a>(
     krate: Crate,
     name: NameToImport,
     assoc_item_search: AssocSearchMode,
-    limit: Option<usize>,
 ) -> impl Iterator<Item = ItemInNs> + 'a {
     let _p = profile::span("items_with_name").detail(|| {
         format!(
-            "Name: {}, crate: {:?}, assoc items: {:?}, limit: {:?}",
+            "Name: {}, crate: {:?}, assoc items: {:?}",
             name.text(),
             assoc_item_search,
             krate.display_name(sema.db).map(|name| name.to_string()),
-            limit,
         )
     });
 
     let prefix = matches!(name, NameToImport::Prefix(..));
-    let (mut local_query, mut external_query) = match name {
+    let (local_query, external_query) = match name {
         NameToImport::Prefix(exact_name, case_sensitive)
         | NameToImport::Exact(exact_name, case_sensitive) => {
             let mut local_query = symbol_index::Query::new(exact_name.clone());
+            local_query.assoc_search_mode(assoc_item_search);
             let mut external_query =
-                // import_map::Query::new(exact_name).assoc_search_mode(assoc_item_search);
+                import_map::Query::new(exact_name).assoc_search_mode(assoc_item_search);
-                import_map::Query::new(exact_name);
             if prefix {
                 local_query.prefix();
                 external_query = external_query.prefix();
@@ -55,6 +53,7 @@ pub fn items_with_name<'a>(
         NameToImport::Fuzzy(fuzzy_search_string, case_sensitive) => {
             let mut local_query = symbol_index::Query::new(fuzzy_search_string.clone());
             local_query.fuzzy();
+            local_query.assoc_search_mode(assoc_item_search);
 
             let mut external_query = import_map::Query::new(fuzzy_search_string.clone())
                 .fuzzy()
@@ -69,18 +68,12 @@ pub fn items_with_name<'a>(
         }
     };
 
-    if let Some(limit) = limit {
+    find_items(sema, krate, local_query, external_query)
-        external_query = external_query.limit(limit);
-        local_query.limit(limit);
-    }
-
-    find_items(sema, krate, assoc_item_search, local_query, external_query)
 }
 
 fn find_items<'a>(
     sema: &'a Semantics<'_, RootDatabase>,
     krate: Crate,
-    assoc_item_search: AssocSearchMode,
     local_query: symbol_index::Query,
     external_query: import_map::Query,
 ) -> impl Iterator<Item = ItemInNs> + 'a {
@@ -98,18 +91,12 @@ fn find_items<'a>(
         });
 
     // Query the local crate using the symbol index.
-    let local_results = local_query
+    let mut local_results = Vec::new();
-        .search(&symbol_index::crate_symbols(db, krate))
+    local_query.search(&symbol_index::crate_symbols(db, krate), |local_candidate| {
-        .into_iter()
+        local_results.push(match local_candidate.def {
-        .filter(move |candidate| match assoc_item_search {
-            AssocSearchMode::Include => true,
-            AssocSearchMode::Exclude => !candidate.is_assoc,
-            AssocSearchMode::AssocItemsOnly => candidate.is_assoc,
-        })
-        .map(|local_candidate| match local_candidate.def {
             hir::ModuleDef::Macro(macro_def) => ItemInNs::Macros(macro_def),
             def => ItemInNs::from(def),
-        });
+        })
-
+    });
-    external_importables.chain(local_results)
+    local_results.into_iter().chain(external_importables)
 }

crates/ide-db/src/symbol_index.rs

@@ -31,9 +31,10 @@ use base_db::{
     salsa::{self, ParallelDatabase},
     SourceDatabaseExt, SourceRootId, Upcast,
 };
-use fst::{self, Streamer};
+use fst::{self, raw::IndexedValue, Automaton, Streamer};
 use hir::{
     db::HirDatabase,
+    import_map::AssocSearchMode,
     symbols::{FileSymbol, SymbolCollector},
     Crate, Module,
 };
@@ -57,8 +58,8 @@ pub struct Query {
     only_types: bool,
     libs: bool,
     mode: SearchMode,
+    assoc_mode: AssocSearchMode,
     case_sensitive: bool,
-    limit: usize,
 }
 
 impl Query {
@@ -70,8 +71,8 @@ impl Query {
             only_types: false,
             libs: false,
             mode: SearchMode::Fuzzy,
+            assoc_mode: AssocSearchMode::Include,
             case_sensitive: false,
-            limit: usize::max_value(),
         }
     }
 
@@ -95,12 +96,13 @@ impl Query {
         self.mode = SearchMode::Prefix;
     }
 
-    pub fn case_sensitive(&mut self) {
+    /// Specifies whether we want to include associated items in the result.
-        self.case_sensitive = true;
+    pub fn assoc_search_mode(&mut self, assoc_mode: AssocSearchMode) {
+        self.assoc_mode = assoc_mode;
     }
 
-    pub fn limit(&mut self, limit: usize) {
+    pub fn case_sensitive(&mut self) {
-        self.limit = limit
+        self.case_sensitive = true;
     }
 }
 
@@ -225,7 +227,9 @@ pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
         indices.iter().flat_map(|indices| indices.iter().cloned()).collect()
     };
 
-    query.search(&indices)
+    let mut res = vec![];
+    query.search(&indices, |f| res.push(f.clone()));
+    res
 }
 
 #[derive(Default)]
@@ -285,6 +289,7 @@ impl SymbolIndex {
             builder.insert(key, value).unwrap();
         }
 
+        // FIXME: fst::Map should ideally have a way to shrink the backing buffer without the unwrap dance
         let map = fst::Map::new({
             let mut buf = builder.into_inner().unwrap();
             buf.shrink_to_fit();
@@ -317,22 +322,54 @@ impl SymbolIndex {
 }
 
 impl Query {
-    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
+    pub(crate) fn search<'sym>(
+        self,
+        indices: &'sym [Arc<SymbolIndex>],
+        cb: impl FnMut(&'sym FileSymbol),
+    ) {
         let _p = profile::span("symbol_index::Query::search");
         let mut op = fst::map::OpBuilder::new();
-        for file_symbols in indices.iter() {
+        match self.mode {
-            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
+            SearchMode::Exact => {
-            op = op.add(file_symbols.map.search(automaton))
+                let automaton = fst::automaton::Str::new(&self.lowercased);
+
+                for index in indices.iter() {
+                    op = op.add(index.map.search(&automaton));
+                }
+                self.search_maps(&indices, op.union(), cb)
+            }
+            SearchMode::Fuzzy => {
+                let automaton = fst::automaton::Subsequence::new(&self.lowercased);
+
+                for index in indices.iter() {
+                    op = op.add(index.map.search(&automaton));
+                }
+                self.search_maps(&indices, op.union(), cb)
+            }
+            SearchMode::Prefix => {
+                let automaton = fst::automaton::Str::new(&self.lowercased).starts_with();
+
+                for index in indices.iter() {
+                    op = op.add(index.map.search(&automaton));
+                }
+                self.search_maps(&indices, op.union(), cb)
+            }
         }
-        let mut stream = op.union();
+    }
-        let mut res = Vec::new();
+
+    fn search_maps<'sym>(
+        &self,
+        indices: &'sym [Arc<SymbolIndex>],
+        mut stream: fst::map::Union<'_>,
+        mut cb: impl FnMut(&'sym FileSymbol),
+    ) {
         while let Some((_, indexed_values)) = stream.next() {
-            for indexed_value in indexed_values {
+            for &IndexedValue { index, value } in indexed_values {
-                let symbol_index = &indices[indexed_value.index];
+                let symbol_index = &indices[index];
-                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);
+                let (start, end) = SymbolIndex::map_value_to_range(value);
 
                 for symbol in &symbol_index.symbols[start..end] {
-                    if self.only_types
+                    let non_type_for_type_only_query = self.only_types
                         && !matches!(
                             symbol.def,
                             hir::ModuleDef::Adt(..)
@@ -340,38 +377,59 @@ impl Query {
                                 | hir::ModuleDef::BuiltinType(..)
                                 | hir::ModuleDef::TraitAlias(..)
                                 | hir::ModuleDef::Trait(..)
-                        )
+                        );
-                    {
+                    if non_type_for_type_only_query || !self.matches_assoc_mode(symbol.is_assoc) {
                         continue;
                     }
-                    let skip = match self.mode {
+                    // FIXME: Deduplicate this from hir-def
+                    let matches = match self.mode {
+                        SearchMode::Exact if self.case_sensitive => symbol.name == self.query,
+                        SearchMode::Exact => symbol.name.eq_ignore_ascii_case(&self.query),
+                        SearchMode::Prefix => {
+                            self.query.len() <= symbol.name.len() && {
+                                let prefix = &symbol.name[..self.query.len() as usize];
+                                if self.case_sensitive {
+                                    prefix == self.query
+                                } else {
+                                    prefix.eq_ignore_ascii_case(&self.query)
+                                }
+                            }
+                        }
                         SearchMode::Fuzzy => {
-                            self.case_sensitive
+                            let mut name = &*symbol.name;
-                                && self.query.chars().any(|c| !symbol.name.contains(c))
+                            self.query.chars().all(|query_char| {
+                                let m = if self.case_sensitive {
+                                    name.match_indices(query_char).next()
+                                } else {
+                                    name.match_indices([
+                                        query_char,
+                                        query_char.to_ascii_uppercase(),
+                                    ])
+                                    .next()
+                                };
+                                match m {
+                                    Some((index, _)) => {
+                                        name = &name[index + 1..];
+                                        true
+                                    }
+                                    None => false,
+                                }
+                            })
                         }
-                        SearchMode::Exact => symbol.name != self.query,
-                        SearchMode::Prefix if self.case_sensitive => {
-                            !symbol.name.starts_with(&self.query)
-                        }
-                        SearchMode::Prefix => symbol
-                            .name
-                            .chars()
-                            .zip(self.lowercased.chars())
-                            .all(|(n, q)| n.to_lowercase().next() == Some(q)),
                     };
-
+                    if matches {
-                    if skip {
+                        cb(symbol);
-                        continue;
-                    }
-
-                    res.push(symbol.clone());
-                    if res.len() >= self.limit {
-                        return res;
                     }
                 }
             }
         }
-        res
+    }
+
+    fn matches_assoc_mode(&self, is_trait_assoc_item: bool) -> bool {
+        match (is_trait_assoc_item, self.assoc_mode) {
+            (true, AssocSearchMode::Exclude) | (false, AssocSearchMode::AssocItemsOnly) => false,
+            _ => true,
+        }
     }
 }
 

crates/ide/src/lib.rs

@@ -414,11 +414,12 @@ impl Analysis {
     }
 
     /// Fuzzy searches for a symbol.
-    pub fn symbol_search(&self, query: Query) -> Cancellable<Vec<NavigationTarget>> {
+    pub fn symbol_search(&self, query: Query, limit: usize) -> Cancellable<Vec<NavigationTarget>> {
         self.with_db(|db| {
             symbol_index::world_symbols(db, query)
                 .into_iter() // xx: should we make this a par iter?
                 .filter_map(|s| s.try_to_nav(db))
+                .take(limit)
                 .map(UpmappingResult::call_site)
                 .collect::<Vec<_>>()
         })

crates/ide/src/navigation_target.rs

@@ -860,7 +860,7 @@ fn foo() { enum FooInner { } }
 "#,
         );
 
-        let navs = analysis.symbol_search(Query::new("FooInner".to_string())).unwrap();
+        let navs = analysis.symbol_search(Query::new("FooInner".to_string()), !0).unwrap();
         expect![[r#"
             [
                 NavigationTarget {
@@ -898,7 +898,7 @@ struct Foo;
 "#,
         );
 
-        let navs = analysis.symbol_search(Query::new("foo".to_string())).unwrap();
+        let navs = analysis.symbol_search(Query::new("foo".to_string()), !0).unwrap();
         assert_eq!(navs.len(), 2)
     }
 }

crates/rust-analyzer/src/handlers/request.rs

@@ -458,7 +458,6 @@ pub(crate) fn handle_workspace_symbol(
 
     let config = snap.config.workspace_symbol();
     let (all_symbols, libs) = decide_search_scope_and_kind(&params, &config);
-    let limit = config.search_limit;
 
     let query = {
         let query: String = params.query.chars().filter(|&c| c != '#' && c != '*').collect();
@@ -469,14 +468,11 @@ pub(crate) fn handle_workspace_symbol(
         if libs {
             q.libs();
         }
-        q.limit(limit);
         q
     };
-    let mut res = exec_query(&snap, query)?;
+    let mut res = exec_query(&snap, query, config.search_limit)?;
     if res.is_empty() && !all_symbols {
-        let mut query = Query::new(params.query);
+        res = exec_query(&snap, Query::new(params.query), config.search_limit)?;
-        query.limit(limit);
-        res = exec_query(&snap, query)?;
     }
 
     return Ok(Some(lsp_types::WorkspaceSymbolResponse::Nested(res)));
@@ -519,9 +515,10 @@ pub(crate) fn handle_workspace_symbol(
     fn exec_query(
         snap: &GlobalStateSnapshot,
         query: Query,
+        limit: usize,
     ) -> anyhow::Result<Vec<lsp_types::WorkspaceSymbol>> {
         let mut res = Vec::new();
-        for nav in snap.analysis.symbol_search(query)? {
+        for nav in snap.analysis.symbol_search(query, limit)? {
             let container_name = nav.container_name.as_ref().map(|v| v.to_string());
 
             let info = lsp_types::WorkspaceSymbol {