Auto merge of rust-lang#16270 - Veykril:flyimport-perf, r=Veykril

internal: Speed up import searching some more

Pushes the sorting to the caller, meaning additional filtering can be done pre-sorting. Similarly a collect call was pushed to the caller for allowing some other filters to run pre-collecting.

Author: bors

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

@@ -89,12 +89,14 @@ use crate::{AssistContext, AssistId, AssistKind, Assists, GroupLabel};
 // ```
 pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
     let (import_assets, syntax_under_caret) = find_importable_node(ctx)?;
-    let mut proposed_imports = import_assets.search_for_imports(
-        &ctx.sema,
-        ctx.config.insert_use.prefix_kind,
-        ctx.config.prefer_no_std,
-        ctx.config.prefer_no_std,
-    );
+    let mut proposed_imports: Vec<_> = import_assets
+        .search_for_imports(
+            &ctx.sema,
+            ctx.config.insert_use.prefix_kind,
+            ctx.config.prefer_no_std,
+            ctx.config.prefer_no_std,
+        )
+        .collect();
     if proposed_imports.is_empty() {
         return None;
     }
@@ -113,6 +115,7 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
     )?;
 
     // we aren't interested in different namespaces
+    proposed_imports.sort_by(|a, b| a.import_path.cmp(&b.import_path));
     proposed_imports.dedup_by(|a, b| a.import_path == b.import_path);
 
     let current_node = match ctx.covering_element() {

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

@@ -37,11 +37,9 @@ use crate::{
 // ```
 pub(crate) fn qualify_path(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
     let (import_assets, syntax_under_caret) = find_importable_node(ctx)?;
-    let mut proposed_imports = import_assets.search_for_relative_paths(
-        &ctx.sema,
-        ctx.config.prefer_no_std,
-        ctx.config.prefer_prelude,
-    );
+    let mut proposed_imports: Vec<_> = import_assets
+        .search_for_relative_paths(&ctx.sema, ctx.config.prefer_no_std, ctx.config.prefer_prelude)
+        .collect();
     if proposed_imports.is_empty() {
         return None;
     }
@@ -82,6 +80,7 @@ pub(crate) fn qualify_path(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option
     };
 
     // we aren't interested in different namespaces
+    proposed_imports.sort_by(|a, b| a.import_path.cmp(&b.import_path));
     proposed_imports.dedup_by(|a, b| a.import_path == b.import_path);
 
     let group_label = group_label(candidate);

crates/ide-completion/src/completions/flyimport.rs

@@ -263,16 +263,21 @@ fn import_on_the_fly(
             ctx.config.prefer_no_std,
             ctx.config.prefer_prelude,
         )
-        .into_iter()
         .filter(ns_filter)
         .filter(|import| {
             let original_item = &import.original_item;
             !ctx.is_item_hidden(&import.item_to_import)
                 && !ctx.is_item_hidden(original_item)
                 && ctx.check_stability(original_item.attrs(ctx.db).as_deref())
         })
-        .sorted_by_key(|located_import| {
-            compute_fuzzy_completion_order_key(&located_import.import_path, &user_input_lowercased)
+        .sorted_by(|a, b| {
+            let key = |import_path| {
+                (
+                    compute_fuzzy_completion_order_key(import_path, &user_input_lowercased),
+                    import_path,
+                )
+            };
+            key(&a.import_path).cmp(&key(&b.import_path))
         })
         .filter_map(|import| {
             render_resolution_with_import(RenderContext::new(ctx), path_ctx, import)
@@ -310,16 +315,21 @@ fn import_on_the_fly_pat_(
             ctx.config.prefer_no_std,
             ctx.config.prefer_prelude,
         )
-        .into_iter()
         .filter(ns_filter)
         .filter(|import| {
             let original_item = &import.original_item;
             !ctx.is_item_hidden(&import.item_to_import)
                 && !ctx.is_item_hidden(original_item)
                 && ctx.check_stability(original_item.attrs(ctx.db).as_deref())
         })
-        .sorted_by_key(|located_import| {
-            compute_fuzzy_completion_order_key(&located_import.import_path, &user_input_lowercased)
+        .sorted_by(|a, b| {
+            let key = |import_path| {
+                (
+                    compute_fuzzy_completion_order_key(import_path, &user_input_lowercased),
+                    import_path,
+                )
+            };
+            key(&a.import_path).cmp(&key(&b.import_path))
         })
         .filter_map(|import| {
             render_resolution_with_import_pat(RenderContext::new(ctx), pattern_ctx, import)
@@ -352,13 +362,18 @@ fn import_on_the_fly_method(
             ctx.config.prefer_no_std,
             ctx.config.prefer_prelude,
         )
-        .into_iter()
         .filter(|import| {
             !ctx.is_item_hidden(&import.item_to_import)
                 && !ctx.is_item_hidden(&import.original_item)
         })
-        .sorted_by_key(|located_import| {
-            compute_fuzzy_completion_order_key(&located_import.import_path, &user_input_lowercased)
+        .sorted_by(|a, b| {
+            let key = |import_path| {
+                (
+                    compute_fuzzy_completion_order_key(import_path, &user_input_lowercased),
+                    import_path,
+                )
+            };
+            key(&a.import_path).cmp(&key(&b.import_path))
         })
         .for_each(|import| match import.original_item {
             ItemInNs::Values(hir::ModuleDef::Function(f)) => {
@@ -407,7 +422,8 @@ fn compute_fuzzy_completion_order_key(
 ) -> usize {
     cov_mark::hit!(certain_fuzzy_order_test);
     let import_name = match proposed_mod_path.segments().last() {
-        Some(name) => name.to_smol_str().to_lowercase(),
+        // FIXME: nasty alloc, this is a hot path!
+        Some(name) => name.to_smol_str().to_ascii_lowercase(),
         None => return usize::MAX,
     };
     match import_name.match_indices(user_input_lowercased).next() {

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

@@ -207,7 +207,7 @@ impl ImportAssets {
         prefix_kind: PrefixKind,
         prefer_no_std: bool,
         prefer_prelude: bool,
-    ) -> Vec<LocatedImport> {
+    ) -> impl Iterator<Item = LocatedImport> {
         let _p = profile::span("import_assets::search_for_imports");
         self.search_for(sema, Some(prefix_kind), prefer_no_std, prefer_prelude)
     }
@@ -218,7 +218,7 @@ impl ImportAssets {
         sema: &Semantics<'_, RootDatabase>,
         prefer_no_std: bool,
         prefer_prelude: bool,
-    ) -> Vec<LocatedImport> {
+    ) -> impl Iterator<Item = LocatedImport> {
         let _p = profile::span("import_assets::search_for_relative_paths");
         self.search_for(sema, None, prefer_no_std, prefer_prelude)
     }
@@ -259,9 +259,15 @@ impl ImportAssets {
         prefixed: Option<PrefixKind>,
         prefer_no_std: bool,
         prefer_prelude: bool,
-    ) -> Vec<LocatedImport> {
+    ) -> impl Iterator<Item = LocatedImport> {
         let _p = profile::span("import_assets::search_for");
 
+        let scope = match sema.scope(&self.candidate_node) {
+            Some(it) => it,
+            None => return <FxHashSet<_>>::default().into_iter(),
+        };
+
+        let krate = self.module_with_candidate.krate();
         let scope_definitions = self.scope_definitions(sema);
         let mod_path = |item| {
             get_mod_path(
@@ -272,30 +278,30 @@ impl ImportAssets {
                 prefer_no_std,
                 prefer_prelude,
             )
-        };
-
-        let krate = self.module_with_candidate.krate();
-        let scope = match sema.scope(&self.candidate_node) {
-            Some(it) => it,
-            None => return Vec::new(),
+            .filter(|path| path.len() > 1)
         };
 
         match &self.import_candidate {
             ImportCandidate::Path(path_candidate) => {
-                path_applicable_imports(sema, krate, path_candidate, mod_path)
-            }
-            ImportCandidate::TraitAssocItem(trait_candidate) => {
-                trait_applicable_items(sema, krate, &scope, trait_candidate, true, mod_path)
-            }
-            ImportCandidate::TraitMethod(trait_candidate) => {
-                trait_applicable_items(sema, krate, &scope, trait_candidate, false, mod_path)
+                path_applicable_imports(sema, krate, path_candidate, mod_path, |item_to_import| {
+                    !scope_definitions.contains(&ScopeDef::from(item_to_import))
+                })
             }
+            ImportCandidate::TraitAssocItem(trait_candidate)
+            | ImportCandidate::TraitMethod(trait_candidate) => trait_applicable_items(
+                sema,
+                krate,
+                &scope,
+                trait_candidate,
+                matches!(self.import_candidate, ImportCandidate::TraitAssocItem(_)),
+                mod_path,
+                |trait_to_import| {
+                    !scope_definitions
+                        .contains(&ScopeDef::ModuleDef(ModuleDef::Trait(trait_to_import)))
+                },
+            ),
         }
         .into_iter()
-        .filter(|import| import.import_path.len() > 1)
-        .filter(|import| !scope_definitions.contains(&ScopeDef::from(import.item_to_import)))
-        .sorted_by(|a, b| a.import_path.cmp(&b.import_path))
-        .collect()
     }
 
     fn scope_definitions(&self, sema: &Semantics<'_, RootDatabase>) -> FxHashSet<ScopeDef> {
@@ -315,6 +321,7 @@ fn path_applicable_imports(
     current_crate: Crate,
     path_candidate: &PathImportCandidate,
     mod_path: impl Fn(ItemInNs) -> Option<ModPath> + Copy,
+    scope_filter: impl Fn(ItemInNs) -> bool + Copy,
 ) -> FxHashSet<LocatedImport> {
     let _p = profile::span("import_assets::path_applicable_imports");
 
@@ -335,6 +342,9 @@ fn path_applicable_imports(
                 AssocSearchMode::Exclude,
             )
             .filter_map(|item| {
+                if !scope_filter(item) {
+                    return None;
+                }
                 let mod_path = mod_path(item)?;
                 Some(LocatedImport::new(mod_path, item, item))
             })
@@ -347,7 +357,7 @@ fn path_applicable_imports(
             path_candidate.name.clone(),
             AssocSearchMode::Include,
         )
-        .filter_map(|item| import_for_item(sema.db, mod_path, &qualifier, item))
+        .filter_map(|item| import_for_item(sema.db, mod_path, &qualifier, item, scope_filter))
         .take(DEFAULT_QUERY_SEARCH_LIMIT.inner())
         .collect(),
     }
@@ -358,6 +368,7 @@ fn import_for_item(
     mod_path: impl Fn(ItemInNs) -> Option<ModPath>,
     unresolved_qualifier: &[SmolStr],
     original_item: ItemInNs,
+    scope_filter: impl Fn(ItemInNs) -> bool,
 ) -> Option<LocatedImport> {
     let _p = profile::span("import_assets::import_for_item");
     let [first_segment, ..] = unresolved_qualifier else { return None };
@@ -413,15 +424,16 @@ fn import_for_item(
             // especially in case of lazy completion edit resolutions.
             return None;
         }
-        (false, Some(trait_to_import)) => {
+        (false, Some(trait_to_import)) if scope_filter(trait_to_import) => {
             LocatedImport::new(mod_path(trait_to_import)?, trait_to_import, original_item)
         }
-        (true, None) => {
+        (true, None) if scope_filter(original_item_candidate) => {
             LocatedImport::new(import_path_candidate, original_item_candidate, original_item)
         }
-        (false, None) => {
+        (false, None) if scope_filter(segment_import) => {
             LocatedImport::new(mod_path(segment_import)?, segment_import, original_item)
         }
+        _ => return None,
     })
 }
 
@@ -490,6 +502,7 @@ fn trait_applicable_items(
     trait_candidate: &TraitImportCandidate,
     trait_assoc_item: bool,
     mod_path: impl Fn(ItemInNs) -> Option<ModPath>,
+    scope_filter: impl Fn(hir::Trait) -> bool,
 ) -> FxHashSet<LocatedImport> {
     let _p = profile::span("import_assets::trait_applicable_items");
 
@@ -533,7 +546,8 @@ fn trait_applicable_items(
             None,
             |assoc| {
                 if required_assoc_items.contains(&assoc) {
-                    let located_trait = assoc.containing_trait(db)?;
+                    let located_trait =
+                        assoc.containing_trait(db).filter(|&it| scope_filter(it))?;
                     let trait_item = ItemInNs::from(ModuleDef::from(located_trait));
                     let import_path = trait_import_paths
                         .entry(trait_item)
@@ -558,7 +572,8 @@ fn trait_applicable_items(
             |function| {
                 let assoc = function.as_assoc_item(db)?;
                 if required_assoc_items.contains(&assoc) {
-                    let located_trait = assoc.containing_trait(db)?;
+                    let located_trait =
+                        assoc.containing_trait(db).filter(|&it| scope_filter(it))?;
                     let trait_item = ItemInNs::from(ModuleDef::from(located_trait));
                     let import_path = trait_import_paths
                         .entry(trait_item)