6 files changed, 119 insertions, 86 deletions

diff --git a/crates/assists/src/handlers/replace_derive_with_manual_impl.rs b/crates/assists/src/handlers/replace_derive_with_manual_impl.rs
index 82625516c..453a6cebf 100644
--- a/crates/assists/src/handlers/replace_derive_with_manual_impl.rs
+++ b/crates/assists/src/handlers/replace_derive_with_manual_impl.rs
@@ -62,19 +62,21 @@ 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 = imports_locator::find_imports(&ctx.sema, current_crate, trait_token.text())
-        .into_iter()
-        .filter_map(|candidate: either::Either<hir::ModuleDef, hir::MacroDef>| match candidate {
-            either::Either::Left(hir::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 =
+        imports_locator::find_exact_imports(&ctx.sema, current_crate, trait_token.text())
+            .filter_map(
+                |candidate: either::Either<hir::ModuleDef, hir::MacroDef>| match candidate {
+                    either::Either::Left(hir::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/assists/src/utils/import_assets.rs b/crates/assists/src/utils/import_assets.rs
index f47edbb76..89e62ba79 100644
--- a/crates/assists/src/utils/import_assets.rs
+++ b/crates/assists/src/utils/import_assets.rs
@@ -179,21 +179,25 @@ impl ImportAssets {
             }
         };
 
-        let mut res = imports_locator::find_imports(sema, current_crate, &self.get_search_query())
-            .into_iter()
-            .filter_map(filter)
-            .filter_map(|candidate| {
-                let item: hir::ItemInNs = candidate.either(Into::into, Into::into);
-                if let Some(prefix_kind) = prefixed {
-                    self.module_with_name_to_import.find_use_path_prefixed(db, item, prefix_kind)
-                } else {
-                    self.module_with_name_to_import.find_use_path(db, item)
-                }
-                .map(|path| (path, item))
-            })
-            .filter(|(use_path, _)| !use_path.segments.is_empty())
-            .take(20)
-            .collect::<Vec<_>>();
+        let mut res =
+            imports_locator::find_exact_imports(sema, current_crate, &self.get_search_query())
+                .filter_map(filter)
+                .filter_map(|candidate| {
+                    let item: hir::ItemInNs = candidate.either(Into::into, Into::into);
+                    if let Some(prefix_kind) = prefixed {
+                        self.module_with_name_to_import.find_use_path_prefixed(
+                            db,
+                            item,
+                            prefix_kind,
+                        )
+                    } else {
+                        self.module_with_name_to_import.find_use_path(db, item)
+                    }
+                    .map(|path| (path, item))
+                })
+                .filter(|(use_path, _)| !use_path.segments.is_empty())
+                .take(20)
+                .collect::<Vec<_>>();
         res.sort_by_key(|(path, _)| path.clone());
         res
     }
diff --git a/crates/completion/src/completions/magic.rs b/crates/completion/src/completions/magic.rs
index 34fc35847..272c9a354 100644
--- a/crates/completion/src/completions/magic.rs
+++ b/crates/completion/src/completions/magic.rs
@@ -2,8 +2,8 @@
 
 use assists::utils::{insert_use, mod_path_to_ast, ImportScope};
 use either::Either;
-use hir::{db::HirDatabase, MacroDef, ModuleDef, Query};
-use itertools::Itertools;
+use hir::{db::HirDatabase, MacroDef, ModuleDef};
+use ide_db::imports_locator;
 use syntax::{algo, AstNode};
 use text_edit::TextEdit;
 
@@ -22,42 +22,40 @@ pub(crate) fn complete_magic(acc: &mut Completions, ctx: &CompletionContext) ->
 
     let potential_import_name = ctx.token.to_string();
 
-    let possible_imports = ctx
-        .krate?
-        // TODO kb use imports_locator instead?
-        .query_external_importables(ctx.db, Query::new(&potential_import_name).limit(40))
-        .unique()
-        .filter_map(|import_candidate| {
-            let use_path = match import_candidate {
-                Either::Left(module_def) => current_module.find_use_path(ctx.db, module_def),
-                Either::Right(macro_def) => current_module.find_use_path(ctx.db, macro_def),
-            }?;
-            Some((use_path, additional_completion(ctx.db, import_candidate)))
-        })
-        .filter_map(|(mod_path, additional_completion)| {
-            let mut builder = TextEdit::builder();
-
-            let correct_qualifier = format!(
-                "{}{}",
-                mod_path.segments.last()?,
-                additional_completion.unwrap_or_default()
-            );
-            builder.replace(anchor.syntax().text_range(), correct_qualifier);
-
-            let rewriter = insert_use(&import_scope, mod_path_to_ast(&mod_path), ctx.config.merge);
-            let old_ast = rewriter.rewrite_root()?;
-            algo::diff(&old_ast, &rewriter.rewrite(&old_ast)).into_text_edit(&mut builder);
-
-            let completion_item: CompletionItem = CompletionItem::new(
-                CompletionKind::Magic,
-                ctx.source_range(),
-                mod_path.to_string(),
-            )
-            .kind(CompletionItemKind::Struct)
-            .text_edit(builder.finish())
-            .into();
-            Some(completion_item)
-        });
+    let possible_imports =
+        imports_locator::find_similar_imports(&ctx.sema, ctx.krate?, &potential_import_name)
+            .filter_map(|import_candidate| {
+                let use_path = match import_candidate {
+                    Either::Left(module_def) => current_module.find_use_path(ctx.db, module_def),
+                    Either::Right(macro_def) => current_module.find_use_path(ctx.db, macro_def),
+                }?;
+                Some((use_path, additional_completion(ctx.db, import_candidate)))
+            })
+            .filter_map(|(mod_path, additional_completion)| {
+                let mut builder = TextEdit::builder();
+
+                let correct_qualifier = format!(
+                    "{}{}",
+                    mod_path.segments.last()?,
+                    additional_completion.unwrap_or_default()
+                );
+                builder.replace(anchor.syntax().text_range(), correct_qualifier);
+
+                let rewriter =
+                    insert_use(&import_scope, mod_path_to_ast(&mod_path), ctx.config.merge);
+                let old_ast = rewriter.rewrite_root()?;
+                algo::diff(&old_ast, &rewriter.rewrite(&old_ast)).into_text_edit(&mut builder);
+
+                let completion_item: CompletionItem = CompletionItem::new(
+                    CompletionKind::Magic,
+                    ctx.source_range(),
+                    mod_path.to_string(),
+                )
+                .kind(CompletionItemKind::Struct)
+                .text_edit(builder.finish())
+                .into();
+                Some(completion_item)
+            });
     acc.add_all(possible_imports);
 
     Some(())
diff --git a/crates/completion/src/lib.rs b/crates/completion/src/lib.rs
index e920fa6b5..8323af8b2 100644
--- a/crates/completion/src/lib.rs
+++ b/crates/completion/src/lib.rs
@@ -118,7 +118,7 @@ pub fn completions(
     completions::macro_in_item_position::complete_macro_in_item_position(&mut acc, &ctx);
     completions::trait_impl::complete_trait_impl(&mut acc, &ctx);
     completions::mod_::complete_mod(&mut acc, &ctx);
-    completions::complete_magic::complete_magic(&mut acc, &ctx);
+    completions::magic::complete_magic(&mut acc, &ctx);
 
     Some(acc)
 }
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index ad58a7cfe..4b7ea3aa9 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -49,7 +49,7 @@ pub use hir_def::{
     builtin_type::BuiltinType,
     docs::Documentation,
     find_path::PrefixKind,
-    import_map::Query,
+    import_map::Query as ExternalImportablesQuery,
     item_scope::ItemInNs,
     nameres::ModuleSource,
     path::{ModPath, PathKind},
diff --git a/crates/ide_db/src/imports_locator.rs b/crates/ide_db/src/imports_locator.rs
index e4f4b5427..71fb7207b 100644
--- a/crates/ide_db/src/imports_locator.rs
+++ b/crates/ide_db/src/imports_locator.rs
@@ -1,40 +1,69 @@
 //! This module contains an import search funcionality that is provided to the assists module.
 //! Later, this should be moved away to a separate crate that is accessible from the assists module.
 
-use hir::{Crate, MacroDef, ModuleDef, Query as ImportMapQuery, Semantics};
+use hir::{Crate, ExternalImportablesQuery, MacroDef, ModuleDef, Semantics};
 use syntax::{ast, AstNode, SyntaxKind::NAME};
 
 use crate::{
     defs::{Definition, NameClass},
-    symbol_index::{self, FileSymbol, Query as SymbolQuery},
+    symbol_index::{self, FileSymbol, Query as LocalImportablesQuery},
     RootDatabase,
 };
 use either::Either;
 use rustc_hash::FxHashSet;
 
-pub fn find_imports<'a>(
+pub fn find_exact_imports<'a>(
     sema: &Semantics<'a, RootDatabase>,
     krate: Crate,
     name_to_import: &str,
-) -> Vec<Either<ModuleDef, MacroDef>> {
-    let _p = profile::span("search_for_imports");
+) -> impl Iterator<Item = Either<ModuleDef, MacroDef>> {
+    let _p = profile::span("find_exact_imports");
+    find_imports(
+        sema,
+        krate,
+        {
+            let mut local_query = LocalImportablesQuery::new(name_to_import.to_string());
+            local_query.exact();
+            local_query.limit(40);
+            local_query
+        },
+        ExternalImportablesQuery::new(name_to_import).anchor_end().case_sensitive().limit(40),
+    )
+}
+
+pub fn find_similar_imports<'a>(
+    sema: &Semantics<'a, RootDatabase>,
+    krate: Crate,
+    name_to_import: &str,
+) -> impl Iterator<Item = Either<ModuleDef, MacroDef>> {
+    let _p = profile::span("find_similar_imports");
+    find_imports(
+        sema,
+        krate,
+        {
+            let mut local_query = LocalImportablesQuery::new(name_to_import.to_string());
+            local_query.limit(40);
+            local_query
+        },
+        ExternalImportablesQuery::new(name_to_import).limit(40),
+    )
+}
+
+fn find_imports<'a>(
+    sema: &Semantics<'a, RootDatabase>,
+    krate: Crate,
+    local_query: LocalImportablesQuery,
+    external_query: ExternalImportablesQuery,
+) -> impl Iterator<Item = Either<ModuleDef, MacroDef>> {
+    let _p = profile::span("find_similar_imports");
     let db = sema.db;
 
     // Query dependencies first.
-    let mut candidates: FxHashSet<_> = krate
-        .query_external_importables(
-            db,
-            ImportMapQuery::new(name_to_import).anchor_end().case_sensitive().limit(40),
-        )
-        .collect();
+    let mut candidates: FxHashSet<_> =
+        krate.query_external_importables(db, external_query).collect();
 
     // Query the local crate using the symbol index.
-    let local_results = {
-        let mut query = SymbolQuery::new(name_to_import.to_string());
-        query.exact();
-        query.limit(40);
-        symbol_index::crate_symbols(db, krate.into(), query)
-    };
+    let local_results = symbol_index::crate_symbols(db, krate.into(), local_query);
 
     candidates.extend(
         local_results
@@ -47,7 +76,7 @@ pub fn find_imports<'a>(
             }),
     );
 
-    candidates.into_iter().collect()
+    candidates.into_iter()
 }
 
 fn get_name_definition<'a>(