4 files changed, 200 insertions, 222 deletions

diff --git a/crates/ide_assists/src/handlers/replace_derive_with_manual_impl.rs b/crates/ide_assists/src/handlers/replace_derive_with_manual_impl.rs
index 88fe2fe90..1a98c51ce 100644
--- a/crates/ide_assists/src/handlers/replace_derive_with_manual_impl.rs
+++ b/crates/ide_assists/src/handlers/replace_derive_with_manual_impl.rs
@@ -1,5 +1,5 @@
 use hir::ModuleDef;
-use ide_db::helpers::mod_path_to_ast;
+use ide_db::helpers::{import_assets::NameToImport, mod_path_to_ast};
 use ide_db::items_locator;
 use itertools::Itertools;
 use syntax::{
@@ -65,20 +65,25 @@ pub(crate) fn replace_derive_with_manual_impl(
     let current_module = ctx.sema.scope(annotated_name.syntax()).module()?;
     let current_crate = current_module.krate();
 
-    let found_traits =
-        items_locator::with_exact_name(&ctx.sema, current_crate, trait_token.text().to_string())
-            .into_iter()
-            .filter_map(|item| match ModuleDef::from(item.as_module_def_id()?) {
-                ModuleDef::Trait(trait_) => Some(trait_),
-                _ => None,
-            })
-            .flat_map(|trait_| {
-                current_module
-                    .find_use_path(ctx.sema.db, hir::ModuleDef::Trait(trait_))
-                    .as_ref()
-                    .map(mod_path_to_ast)
-                    .zip(Some(trait_))
-            });
+    let found_traits = items_locator::locate_for_name(
+        &ctx.sema,
+        current_crate,
+        NameToImport::Exact(trait_token.text().to_string()),
+        items_locator::AssocItemSearch::Exclude,
+        None,
+    )
+    .into_iter()
+    .filter_map(|item| match ModuleDef::from(item.as_module_def_id()?) {
+        ModuleDef::Trait(trait_) => Some(trait_),
+        _ => None,
+    })
+    .flat_map(|trait_| {
+        current_module
+            .find_use_path(ctx.sema.db, hir::ModuleDef::Trait(trait_))
+            .as_ref()
+            .map(mod_path_to_ast)
+            .zip(Some(trait_))
+    });
 
     let mut no_traits_found = true;
     for (trait_path, trait_) in found_traits.inspect(|_| no_traits_found = false) {
diff --git a/crates/ide_completion/src/lib.rs b/crates/ide_completion/src/lib.rs
index d9ea7b7ea..c91c98871 100644
--- a/crates/ide_completion/src/lib.rs
+++ b/crates/ide_completion/src/lib.rs
@@ -14,7 +14,10 @@ mod completions;
 use completions::flyimport::position_for_import;
 use ide_db::{
     base_db::FilePosition,
-    helpers::{import_assets::LocatedImport, insert_use::ImportScope},
+    helpers::{
+        import_assets::{LocatedImport, NameToImport},
+        insert_use::ImportScope,
+    },
     items_locator, RootDatabase,
 };
 use text_edit::TextEdit;
@@ -149,15 +152,20 @@ pub fn resolve_completion_edits(
     let current_module = ctx.sema.scope(position_for_import).module()?;
     let current_crate = current_module.krate();
 
-    let (import_path, item_to_import) =
-        items_locator::with_exact_name(&ctx.sema, current_crate, imported_name)
-            .into_iter()
-            .filter_map(|candidate| {
-                current_module
-                    .find_use_path_prefixed(db, candidate, config.insert_use.prefix_kind)
-                    .zip(Some(candidate))
-            })
-            .find(|(mod_path, _)| mod_path.to_string() == full_import_path)?;
+    let (import_path, item_to_import) = items_locator::locate_for_name(
+        &ctx.sema,
+        current_crate,
+        NameToImport::Exact(imported_name),
+        items_locator::AssocItemSearch::Include,
+        None,
+    )
+    .into_iter()
+    .filter_map(|candidate| {
+        current_module
+            .find_use_path_prefixed(db, candidate, config.insert_use.prefix_kind)
+            .zip(Some(candidate))
+    })
+    .find(|(mod_path, _)| mod_path.to_string() == full_import_path)?;
     let import =
         LocatedImport::new(import_path.clone(), item_to_import, item_to_import, Some(import_path));
 
diff --git a/crates/ide_db/src/helpers/import_assets.rs b/crates/ide_db/src/helpers/import_assets.rs
index dbc980ba9..ae234eddc 100644
--- a/crates/ide_db/src/helpers/import_assets.rs
+++ b/crates/ide_db/src/helpers/import_assets.rs
@@ -61,7 +61,7 @@ pub struct FirstSegmentUnresolved {
 }
 
 /// A name that will be used during item lookups.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum NameToImport {
     /// Requires items with names that exactly match the given string, case-sensitive.
     Exact(String),
@@ -201,131 +201,96 @@ impl ImportAssets {
         sema: &Semantics<RootDatabase>,
         prefixed: Option<PrefixKind>,
     ) -> Vec<LocatedImport> {
-        let items_with_candidate_name = match self.name_to_import() {
-            NameToImport::Exact(exact_name) => items_locator::with_exact_name(
-                sema,
-                self.module_with_candidate.krate(),
-                exact_name.clone(),
-            ),
-            // FIXME: ideally, we should avoid using `fst` for seacrhing trait imports for assoc items:
-            // instead, we need to look up all trait impls for a certain struct and search through them only
-            // see https://github.com/rust-analyzer/rust-analyzer/pull/7293#issuecomment-761585032
-            // and https://rust-lang.zulipchat.com/#narrow/stream/185405-t-compiler.2Fwg-rls-2.2E0/topic/Blanket.20trait.20impls.20lookup
-            // for the details
-            NameToImport::Fuzzy(fuzzy_name) => {
-                let (assoc_item_search, limit) = if self.import_candidate.is_trait_candidate() {
-                    (AssocItemSearch::AssocItemsOnly, None)
-                } else {
-                    (AssocItemSearch::Include, Some(DEFAULT_QUERY_SEARCH_LIMIT))
-                };
-
-                items_locator::with_similar_name(
-                    sema,
-                    self.module_with_candidate.krate(),
-                    fuzzy_name.clone(),
-                    assoc_item_search,
-                    limit,
-                )
-            }
-        };
+        let _p = profile::span("import_assets::search_for");
 
         let scope_definitions = self.scope_definitions(sema);
-        self.applicable_defs(sema.db, prefixed, items_with_candidate_name)
-            .into_iter()
-            .filter(|import| import.import_path.len() > 1)
-            .filter(|import| !scope_definitions.contains(&ScopeDef::from(import.item_to_import)))
-            .sorted_by_key(|import| import.import_path.clone())
-            .collect()
-    }
-
-    fn scope_definitions(&self, sema: &Semantics<RootDatabase>) -> FxHashSet<ScopeDef> {
-        let mut scope_definitions = FxHashSet::default();
-        sema.scope(&self.candidate_node).process_all_names(&mut |_, scope_def| {
-            scope_definitions.insert(scope_def);
-        });
-        scope_definitions
-    }
-
-    fn name_to_import(&self) -> &NameToImport {
-        match &self.import_candidate {
-            ImportCandidate::Path(candidate) => &candidate.name,
-            ImportCandidate::TraitAssocItem(candidate)
-            | ImportCandidate::TraitMethod(candidate) => &candidate.assoc_item_name,
-        }
-    }
-
-    fn applicable_defs(
-        &self,
-        db: &RootDatabase,
-        prefixed: Option<PrefixKind>,
-        items_with_candidate_name: FxHashSet<ItemInNs>,
-    ) -> FxHashSet<LocatedImport> {
-        let _p = profile::span("import_assets::applicable_defs");
         let current_crate = self.module_with_candidate.krate();
-
         let mod_path = |item| {
-            get_mod_path(db, item_for_path_search(db, item)?, &self.module_with_candidate, prefixed)
+            get_mod_path(
+                sema.db,
+                item_for_path_search(sema.db, item)?,
+                &self.module_with_candidate,
+                prefixed,
+            )
         };
 
         match &self.import_candidate {
             ImportCandidate::Path(path_candidate) => {
-                path_applicable_imports(db, path_candidate, mod_path, items_with_candidate_name)
+                path_applicable_imports(sema, current_crate, path_candidate, mod_path)
+            }
+            ImportCandidate::TraitAssocItem(trait_candidate) => {
+                trait_applicable_items(sema, current_crate, trait_candidate, true, mod_path)
+            }
+            ImportCandidate::TraitMethod(trait_candidate) => {
+                trait_applicable_items(sema, current_crate, trait_candidate, false, mod_path)
             }
-            ImportCandidate::TraitAssocItem(trait_candidate) => trait_applicable_items(
-                db,
-                current_crate,
-                trait_candidate,
-                true,
-                mod_path,
-                items_with_candidate_name,
-            ),
-            ImportCandidate::TraitMethod(trait_candidate) => trait_applicable_items(
-                db,
-                current_crate,
-                trait_candidate,
-                false,
-                mod_path,
-                items_with_candidate_name,
-            ),
         }
+        .into_iter()
+        .filter(|import| import.import_path.len() > 1)
+        .filter(|import| !scope_definitions.contains(&ScopeDef::from(import.item_to_import)))
+        .sorted_by_key(|import| import.import_path.clone())
+        .collect()
+    }
+
+    fn scope_definitions(&self, sema: &Semantics<RootDatabase>) -> FxHashSet<ScopeDef> {
+        let _p = profile::span("import_assets::scope_definitions");
+        let mut scope_definitions = FxHashSet::default();
+        sema.scope(&self.candidate_node).process_all_names(&mut |_, scope_def| {
+            scope_definitions.insert(scope_def);
+        });
+        scope_definitions
     }
 }
 
 fn path_applicable_imports(
-    db: &RootDatabase,
+    sema: &Semantics<RootDatabase>,
+    current_crate: Crate,
     path_candidate: &PathImportCandidate,
     mod_path: impl Fn(ItemInNs) -> Option<ModPath> + Copy,
-    items_with_candidate_name: FxHashSet<ItemInNs>,
 ) -> FxHashSet<LocatedImport> {
     let _p = profile::span("import_assets::path_applicable_imports");
 
-    let (unresolved_first_segment, unresolved_qualifier) = match &path_candidate.qualifier {
+    match &path_candidate.qualifier {
         None => {
-            return items_with_candidate_name
-                .into_iter()
-                .filter_map(|item| {
-                    if item_as_assoc(db, item).is_some() {
-                        // unqualified assoc items are not valid syntax
-                        return None;
-                    }
-
-                    let mod_path = mod_path(item)?;
-                    Some(LocatedImport::new(mod_path.clone(), item, item, Some(mod_path)))
-                })
-                .collect();
+            items_locator::locate_for_name(
+                sema,
+                current_crate,
+                path_candidate.name.clone(),
+                // unqualified assoc items are not valid syntax
+                AssocItemSearch::Exclude,
+                Some(DEFAULT_QUERY_SEARCH_LIMIT),
+            )
+            .into_iter()
+            .filter_map(|item| {
+                let mod_path = mod_path(item)?;
+                Some(LocatedImport::new(mod_path.clone(), item, item, Some(mod_path)))
+            })
+            .collect()
         }
-        Some(first_segment_unresolved) => (
-            first_segment_unresolved.fist_segment.to_string(),
-            path_to_string_stripping_turbo_fish(&first_segment_unresolved.full_qualifier),
-        ),
-    };
-
-    items_with_candidate_name
-        .into_iter()
-        .filter_map(|item| {
-            import_for_item(db, mod_path, &unresolved_first_segment, &unresolved_qualifier, item)
-        })
-        .collect()
+        Some(first_segment_unresolved) => {
+            let unresolved_qualifier =
+                path_to_string_stripping_turbo_fish(&first_segment_unresolved.full_qualifier);
+            let unresolved_first_segment = first_segment_unresolved.fist_segment.text();
+            items_locator::locate_for_name(
+                sema,
+                current_crate,
+                path_candidate.name.clone(),
+                AssocItemSearch::Include,
+                Some(DEFAULT_QUERY_SEARCH_LIMIT),
+            )
+            .into_iter()
+            .filter_map(|item| {
+                import_for_item(
+                    sema.db,
+                    mod_path,
+                    unresolved_first_segment,
+                    &unresolved_qualifier,
+                    item,
+                )
+            })
+            .collect()
+        }
+    }
 }
 
 fn import_for_item(
@@ -440,25 +405,32 @@ fn module_with_segment_name(
 }
 
 fn trait_applicable_items(
-    db: &RootDatabase,
+    sema: &Semantics<RootDatabase>,
     current_crate: Crate,
     trait_candidate: &TraitImportCandidate,
     trait_assoc_item: bool,
     mod_path: impl Fn(ItemInNs) -> Option<ModPath>,
-    items_with_candidate_name: FxHashSet<ItemInNs>,
 ) -> FxHashSet<LocatedImport> {
     let _p = profile::span("import_assets::trait_applicable_items");
-    let mut required_assoc_items = FxHashSet::default();
 
-    let trait_candidates = items_with_candidate_name
-        .into_iter()
-        .filter_map(|input| item_as_assoc(db, input))
-        .filter_map(|assoc| {
-            let assoc_item_trait = assoc.containing_trait(db)?;
-            required_assoc_items.insert(assoc);
-            Some(assoc_item_trait.into())
-        })
-        .collect();
+    let db = sema.db;
+
+    let mut required_assoc_items = FxHashSet::default();
+    let trait_candidates = items_locator::locate_for_name(
+        sema,
+        current_crate,
+        trait_candidate.assoc_item_name.clone(),
+        AssocItemSearch::AssocItemsOnly,
+        Some(DEFAULT_QUERY_SEARCH_LIMIT),
+    )
+    .into_iter()
+    .filter_map(|input| item_as_assoc(db, input))
+    .filter_map(|assoc| {
+        let assoc_item_trait = assoc.containing_trait(db)?;
+        required_assoc_items.insert(assoc);
+        Some(assoc_item_trait.into())
+    })
+    .collect();
 
     let mut located_imports = FxHashSet::default();
 
@@ -567,10 +539,6 @@ impl ImportCandidate {
     ) -> Option<Self> {
         path_import_candidate(sema, qualifier, NameToImport::Fuzzy(fuzzy_name))
     }
-
-    fn is_trait_candidate(&self) -> bool {
-        matches!(self, ImportCandidate::TraitAssocItem(_) | ImportCandidate::TraitMethod(_))
-    }
 }
 
 fn path_import_candidate(
diff --git a/crates/ide_db/src/items_locator.rs b/crates/ide_db/src/items_locator.rs
index 8a7f02935..088be72c4 100644
--- a/crates/ide_db/src/items_locator.rs
+++ b/crates/ide_db/src/items_locator.rs
@@ -10,6 +10,7 @@ use syntax::{ast, AstNode, SyntaxKind::NAME};
 
 use crate::{
     defs::{Definition, NameClass},
+    helpers::import_assets::NameToImport,
     symbol_index::{self, FileSymbol},
     RootDatabase,
 };
@@ -17,115 +18,105 @@ use rustc_hash::FxHashSet;
 
 pub(crate) const DEFAULT_QUERY_SEARCH_LIMIT: usize = 40;
 
-pub fn with_exact_name(
-    sema: &Semantics<'_, RootDatabase>,
-    krate: Crate,
-    exact_name: String,
-) -> FxHashSet<ItemInNs> {
-    let _p = profile::span("find_exact_imports");
-    find_items(
-        sema,
-        krate,
-        {
-            let mut local_query = symbol_index::Query::new(exact_name.clone());
-            local_query.exact();
-            local_query.limit(DEFAULT_QUERY_SEARCH_LIMIT);
-            local_query
-        },
-        import_map::Query::new(exact_name)
-            .limit(DEFAULT_QUERY_SEARCH_LIMIT)
-            .name_only()
-            .search_mode(import_map::SearchMode::Equals)
-            .case_sensitive(),
-    )
-}
-
-#[derive(Debug)]
+/// TODO kb docs here and around + update the module doc
+#[derive(Debug, Clone, Copy)]
 pub enum AssocItemSearch {
     Include,
     Exclude,
     AssocItemsOnly,
 }
 
-pub fn with_similar_name(
+pub fn locate_for_name(
     sema: &Semantics<'_, RootDatabase>,
     krate: Crate,
-    fuzzy_search_string: String,
+    name: NameToImport,
     assoc_item_search: AssocItemSearch,
     limit: Option<usize>,
 ) -> FxHashSet<ItemInNs> {
-    let _p = profile::span("find_similar_imports");
+    let _p = profile::span("locate_for_name").detail(|| {
+        format!(
+            "Name: {} ({:?}), crate: {:?}, limit: {:?}",
+            name.text(),
+            assoc_item_search,
+            krate.display_name(sema.db).map(|name| name.to_string()),
+            limit,
+        )
+    });
+
+    let (mut local_query, mut external_query) = match name {
+        NameToImport::Exact(exact_name) => {
+            let mut local_query = symbol_index::Query::new(exact_name.clone());
+            local_query.exact();
 
-    let mut external_query = import_map::Query::new(fuzzy_search_string.clone())
-        .search_mode(import_map::SearchMode::Fuzzy)
-        .name_only();
+            let external_query = import_map::Query::new(exact_name)
+                .name_only()
+                .search_mode(import_map::SearchMode::Equals)
+                .case_sensitive();
 
-    match assoc_item_search {
-        AssocItemSearch::Include => {}
-        AssocItemSearch::Exclude => {
-            external_query = external_query.exclude_import_kind(ImportKind::AssociatedItem);
+            (local_query, external_query)
         }
-        AssocItemSearch::AssocItemsOnly => {
-            external_query = external_query.assoc_items_only();
+        NameToImport::Fuzzy(fuzzy_search_string) => {
+            let mut external_query = import_map::Query::new(fuzzy_search_string.clone())
+                .search_mode(import_map::SearchMode::Fuzzy)
+                .name_only();
+            match assoc_item_search {
+                AssocItemSearch::Include => {}
+                AssocItemSearch::Exclude => {
+                    external_query = external_query.exclude_import_kind(ImportKind::AssociatedItem);
+                }
+                AssocItemSearch::AssocItemsOnly => {
+                    external_query = external_query.assoc_items_only();
+                }
+            }
+
+            (symbol_index::Query::new(fuzzy_search_string), external_query)
         }
-    }
-
-    let mut local_query = symbol_index::Query::new(fuzzy_search_string);
+    };
 
     if let Some(limit) = limit {
         external_query = external_query.limit(limit);
         local_query.limit(limit);
     }
 
-    find_items(sema, krate, local_query, external_query)
-        .into_iter()
-        .filter(move |&item| match assoc_item_search {
-            AssocItemSearch::Include => true,
-            AssocItemSearch::Exclude => !is_assoc_item(item, sema.db),
-            AssocItemSearch::AssocItemsOnly => is_assoc_item(item, sema.db),
-        })
-        .collect()
-}
-
-fn is_assoc_item(item: ItemInNs, db: &RootDatabase) -> bool {
-    item.as_module_def_id()
-        .and_then(|module_def_id| ModuleDef::from(module_def_id).as_assoc_item(db))
-        .is_some()
+    find_items(sema, krate, assoc_item_search, local_query, external_query)
 }
 
 fn find_items(
     sema: &Semantics<'_, RootDatabase>,
     krate: Crate,
+    assoc_item_search: AssocItemSearch,
     local_query: symbol_index::Query,
     external_query: import_map::Query,
 ) -> FxHashSet<ItemInNs> {
-    let _p = profile::span("find_similar_imports");
+    let _p = profile::span("find_items");
     let db = sema.db;
 
-    // Query dependencies first.
-    let mut candidates = krate
-        .query_external_importables(db, external_query)
-        .map(|external_importable| match external_importable {
-            Either::Left(module_def) => ItemInNs::from(module_def),
-            Either::Right(macro_def) => ItemInNs::from(macro_def),
-        })
-        .collect::<FxHashSet<_>>();
+    let external_importables =
+        krate.query_external_importables(db, external_query).map(|external_importable| {
+            match external_importable {
+                Either::Left(module_def) => ItemInNs::from(module_def),
+                Either::Right(macro_def) => ItemInNs::from(macro_def),
+            }
+        });
 
     // Query the local crate using the symbol index.
-    let local_results = symbol_index::crate_symbols(db, krate.into(), local_query);
-
-    candidates.extend(
-        local_results
-            .into_iter()
-            .filter_map(|local_candidate| get_name_definition(sema, &local_candidate))
-            .filter_map(|name_definition_to_import| match name_definition_to_import {
-                Definition::ModuleDef(module_def) => Some(ItemInNs::from(module_def)),
-                Definition::Macro(macro_def) => Some(ItemInNs::from(macro_def)),
-                _ => None,
-            }),
-    );
-
-    candidates
+    let local_results = symbol_index::crate_symbols(db, krate.into(), local_query)
+        .into_iter()
+        .filter_map(|local_candidate| get_name_definition(sema, &local_candidate))
+        .filter_map(|name_definition_to_import| match name_definition_to_import {
+            Definition::ModuleDef(module_def) => Some(ItemInNs::from(module_def)),
+            Definition::Macro(macro_def) => Some(ItemInNs::from(macro_def)),
+            _ => None,
+        });
+
+    external_importables
+        .chain(local_results)
+        .filter(move |&item| match assoc_item_search {
+            AssocItemSearch::Include => true,
+            AssocItemSearch::Exclude => !is_assoc_item(item, sema.db),
+            AssocItemSearch::AssocItemsOnly => is_assoc_item(item, sema.db),
+        })
+        .collect()
 }
 
 fn get_name_definition(
@@ -144,3 +135,9 @@ fn get_name_definition(
     let name = ast::Name::cast(candidate_name_node)?;
     NameClass::classify(sema, &name)?.defined(sema.db)
 }
+
+fn is_assoc_item(item: ItemInNs, db: &RootDatabase) -> bool {
+    item.as_module_def_id()
+        .and_then(|module_def_id| ModuleDef::from(module_def_id).as_assoc_item(db))
+        .is_some()
+}