WIP - Sort suggested imports by type for data types

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

@@ -135,6 +135,7 @@ pub(crate) fn import_on_the_fly_path(
         Qualified::With { path, .. } => Some(path.clone()),
         _ => None,
     };
+
     let import_assets = import_assets_for_path(ctx, &potential_import_name, qualifier.clone())?;
 
     import_on_the_fly(
@@ -257,24 +258,42 @@ fn import_on_the_fly(
     };
     let user_input_lowercased = potential_import_name.to_lowercase();
 
-    import_assets
-        .search_for_imports(&ctx.sema, ctx.config.insert_use.prefix_kind, ctx.config.prefer_no_std)
-        .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)
-        })
-        .filter_map(|import| {
-            render_resolution_with_import(RenderContext::new(ctx), path_ctx, import)
-        })
-        .map(|builder| builder.build(ctx.db))
-        .for_each(|item| acc.add(item));
+    let (expected_ty_imports, other_imports): (Vec<LocatedImport>, Vec<LocatedImport>) =
+        import_assets
+            .search_for_imports(
+                &ctx.sema,
+                ctx.config.insert_use.prefix_kind,
+                ctx.config.prefer_no_std,
+            )
+            .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())
+            })
+            .partition(|located_import| is_expected_type(ctx, located_import));
+
+    // Note that the ordering of `expected_ty_imports` and `other_imports` matters.
+    // Imports with the expected type should be displayed before other imports and must
+    // therefore be sorted separately and added to the Completions first.
+    for imports in [expected_ty_imports, other_imports] {
+        imports
+            .into_iter()
+            .sorted_by_key(|located_import| {
+                compute_fuzzy_completion_order_key(
+                    &located_import.import_path,
+                    &user_input_lowercased,
+                )
+            })
+            .filter_map(|import| {
+                render_resolution_with_import(RenderContext::new(ctx), path_ctx, import)
+            })
+            .map(|builder| builder.build(ctx.db))
+            .for_each(|item| acc.add(item));
+    }
+
     Some(())
 }
 
@@ -377,13 +396,29 @@ fn import_assets_for_path(
         &ctx.sema,
         ctx.token.parent()?,
     )?;
+
     if fuzzy_name_length < 3 {
         cov_mark::hit!(flyimport_exact_on_short_path);
         assets_for_path.path_fuzzy_name_to_exact(false);
     }
+
     Some(assets_for_path)
 }
 
+fn is_expected_type(ctx: &CompletionContext<'_>, located_import: &LocatedImport) -> bool {
+    let Some(expected_type): Option<ModuleDef> =
+        ctx.expected_type.as_ref().map(|et| et.as_adt()).flatten().map(Into::into)
+    else {
+        return false;
+    };
+
+    let Some(import_type) = located_import.original_item.as_module_def() else {
+        return false;
+    };
+
+    expected_type == import_type
+}
+
 fn compute_fuzzy_completion_order_key(
     proposed_mod_path: &hir::ModPath,
     user_input_lowercased: &str,
@@ -393,6 +428,7 @@ fn compute_fuzzy_completion_order_key(
         Some(name) => name.to_smol_str().to_lowercase(),
         None => return usize::MAX,
     };
+
     match import_name.match_indices(user_input_lowercased).next() {
         Some((first_matching_index, _)) => first_matching_index,
         None => usize::MAX,

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

@@ -389,6 +389,87 @@ fn main() {
     );
 }
 
+#[test]
+fn sort_suggested_structs_by_types() {
+    check(
+        r#"
+//- /lib.rs crate:dep
+pub mod test_mod_a {
+    pub struct Struct {}
+}
+
+pub mod test_mod_b {
+    pub struct Struct {}
+}
+
+//- /main.rs crate:main deps:dep
+
+fn test(input: dep::test_mod_b::Struct) { }
+
+fn main() {
+    test(Struct$0);
+}
+"#,
+        expect![[r#"st Struct (use dep::test_mod_b::Struct)
+st Struct (use dep::test_mod_a::Struct)
+"#]],
+    );
+}
+
+#[test]
+fn sort_suggested_enums_by_types() {
+    check(
+        r#"
+//- /lib.rs crate:dep
+pub mod test_mod_a {
+    pub enum Enum {}
+}
+
+pub mod test_mod_b {
+    pub enum Enum {}
+}
+
+//- /main.rs crate:main deps:dep
+
+fn test(input: dep::test_mod_b::Enum) {  }
+
+fn main() {
+    test(Enum$0);
+}
+"#,
+        expect![[r#"en Enum (use dep::test_mod_b::Enum)
+en Enum (use dep::test_mod_a::Enum)
+"#]],
+    );
+}
+
+#[test]
+fn sort_suggested_unions_by_types() {
+    check(
+        r#"
+//- /lib.rs crate:dep
+pub mod test_mod_a {
+    pub union Union {}
+}
+
+pub mod test_mod_b {
+    pub union Union {}
+}
+
+//- /main.rs crate:main deps:dep
+
+fn test(input: dep::test_mod_b::Union) {  }
+
+fn main() {
+    test(Union$0);
+}
+"#,
+        expect![[r#"un Union (use dep::test_mod_b::Union)
+un Union (use dep::test_mod_a::Union)
+"#]],
+    );
+}
+
 #[test]
 fn does_not_propose_names_in_scope() {
     check(