Fix symbol_index::Query::search_maps not adhering to case-sensitivity correctly

crates/hir-def/src/import_map.rs

@@ -406,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)
                 })
@@ -417,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

crates/hir/src/symbols.rs

@@ -163,7 +163,7 @@ fn collect_from_module(&mut self, module_id: ModuleId) {
         }
 
         // 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-db/src/items_locator.rs

@@ -36,9 +36,9 @@ pub fn items_with_name<'a>(
         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 +55,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()
@@ -73,13 +74,12 @@ pub fn items_with_name<'a>(
         local_query.limit(limit);
     }
 
-    find_items(sema, krate, assoc_item_search, local_query, external_query)
+    find_items(sema, krate, 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 {
@@ -97,18 +97,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),
+        })
     });
-
+    local_results.into_iter().chain(external_importables)
-    external_importables.chain(local_results)
 }

crates/ide-db/src/symbol_index.rs

@@ -31,9 +31,10 @@
     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,6 +58,7 @@ pub struct Query {
     only_types: bool,
     libs: bool,
     mode: SearchMode,
+    assoc_mode: AssocSearchMode,
     case_sensitive: bool,
     limit: usize,
 }
@@ -70,6 +72,7 @@ pub fn new(query: String) -> Query {
             only_types: false,
             libs: false,
             mode: SearchMode::Fuzzy,
+            assoc_mode: AssocSearchMode::Include,
             case_sensitive: false,
             limit: usize::max_value(),
         }
@@ -95,6 +98,11 @@ pub fn prefix(&mut self) {
         self.mode = SearchMode::Prefix;
     }
 
+    /// Specifies whether we want to include associated items in the result.
+    pub fn assoc_search_mode(&mut self, assoc_mode: AssocSearchMode) {
+        self.assoc_mode = assoc_mode;
+    }
+
     pub fn case_sensitive(&mut self) {
         self.case_sensitive = true;
     }
@@ -225,7 +233,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 +295,7 @@ fn cmp(lhs: &FileSymbol, rhs: &FileSymbol) -> Ordering {
             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 +328,54 @@ fn map_value_to_range(value: u64) -> (usize, usize) {
 }
 
 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));
                 }
-        let mut stream = op.union();
+                self.search_maps(&indices, op.union(), cb)
-        let mut res = Vec::new();
+            }
+            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)
+            }
+        }
+    }
+
+    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 +383,59 @@ pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
                                 | 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 {
-                        SearchMode::Exact => symbol.name != self.query,
+                                    name.match_indices(query_char).next()
-                        SearchMode::Prefix if self.case_sensitive => {
+                                } else {
-                            !symbol.name.starts_with(&self.query)
+                                    name.match_indices([
-                        }
+                                        query_char,
-                        SearchMode::Prefix => symbol
+                                        query_char.to_ascii_uppercase(),
-                            .name
+                                    ])
-                            .chars()
+                                    .next()
-                            .zip(self.lowercased.chars())
-                            .all(|(n, q)| n.to_lowercase().next() == Some(q)),
                                 };
-
+                                match m {
-                    if skip {
+                                    Some((index, _)) => {
-                        continue;
+                                        name = &name[index + 1..];
+                                        true
+                                    }
+                                    None => false,
+                                }
+                            })
+                        }
+                    };
+                    if matches {
+                        cb(symbol);
                     }
-
-                    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,
+        }
     }
 }