feat: Allow crate authors to declare that their trait prefers to be imported as _

crates/hir-def/src/attrs.rs

@@ -258,6 +258,9 @@ fn match_attr_flags(attr_flags: &mut AttrFlags, attr: ast::Meta) -> ControlFlow<
                 Some(second_segment) => match &*first_segment {
                     "rust_analyzer" => match &*second_segment {
                         "skip" => attr_flags.insert(AttrFlags::RUST_ANALYZER_SKIP),
+                        "prefer_underscore_import" => {
+                            attr_flags.insert(AttrFlags::PREFER_UNDERSCORE_IMPORT)
+                        }
                         _ => {}
                     },
                     _ => {}
@@ -330,6 +333,8 @@ bitflags::bitflags! {
         const MACRO_STYLE_BRACES = 1 << 46;
         const MACRO_STYLE_BRACKETS = 1 << 47;
         const MACRO_STYLE_PARENTHESES = 1 << 48;
+
+        const PREFER_UNDERSCORE_IMPORT = 1 << 49;
     }
 }
 

crates/hir/src/lib.rs

@@ -3319,6 +3319,19 @@ impl Trait {
     pub fn complete(self, db: &dyn HirDatabase) -> Complete {
         Complete::extract(true, self.attrs(db).attrs)
     }
+
+    // Feature: Prefer Underscore Import Attribute
+    // Crate authors can declare that their trait prefers to be imported `as _`. This can be used
+    // for example for extension traits. To do that, a trait has to include the attribute
+    // `#[rust_analyzer::prefer_underscore_import]`
+    //
+    // When a trait includes this attribute, flyimport will import it `as _`, and the quickfix
+    // to import it will prefer to import it `as _` (but allow to import it normally as well).
+    //
+    // Malformed attributes will be ignored without warnings.
+    pub fn prefer_underscore_import(self, db: &dyn HirDatabase) -> bool {
+        AttrFlags::query(db, self.id.into()).contains(AttrFlags::PREFER_UNDERSCORE_IMPORT)
+    }
 }
 
 impl HasVisibility for Trait {

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

@@ -6,7 +6,7 @@ use ide_db::{
     active_parameter::ActiveParameter,
     helpers::mod_path_to_ast,
     imports::{
-        import_assets::{ImportAssets, ImportCandidate, LocatedImport},
+        import_assets::{ImportAssets, ImportCandidate, LocatedImport, TraitImportCandidate},
         insert_use::{ImportScope, insert_use, insert_use_as_alias},
     },
 };
@@ -123,44 +123,48 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
 
         let (assist_id, import_name) =
             (AssistId::quick_fix("auto_import"), import_path.display(ctx.db(), edition));
-        acc.add_group(
-            &group_label,
-            assist_id,
-            format!("Import `{import_name}`"),
-            range,
-            |builder| {
+        let add_normal_import = |acc: &mut Assists, label| {
+            acc.add_group(&group_label, assist_id, label, range, |builder| {
                 let scope = builder.make_import_scope_mut(scope.clone());
                 insert_use(&scope, mod_path_to_ast(&import_path, edition), &ctx.config.insert_use);
-            },
-        );
-
-        match import_assets.import_candidate() {
-            ImportCandidate::TraitAssocItem(name) | ImportCandidate::TraitMethod(name) => {
-                let is_method =
-                    matches!(import_assets.import_candidate(), ImportCandidate::TraitMethod(_));
-                let type_ = if is_method { "method" } else { "item" };
-                let group_label = GroupLabel(format!(
-                    "Import a trait for {} {} by alias",
-                    type_,
-                    name.assoc_item_name.text()
-                ));
-                acc.add_group(
-                    &group_label,
-                    assist_id,
-                    format!("Import `{import_name} as _`"),
-                    range,
-                    |builder| {
-                        let scope = builder.make_import_scope_mut(scope.clone());
-                        insert_use_as_alias(
-                            &scope,
-                            mod_path_to_ast(&import_path, edition),
-                            &ctx.config.insert_use,
-                            edition,
-                        );
-                    },
+            })
+        };
+        let add_underscore_import = |acc: &mut Assists, name: &TraitImportCandidate<'_>, label| {
+            let is_method =
+                matches!(import_assets.import_candidate(), ImportCandidate::TraitMethod(_));
+            let type_ = if is_method { "method" } else { "item" };
+            let group_label = GroupLabel(format!(
+                "Import a trait for {} {} by alias",
+                type_,
+                name.assoc_item_name.text()
+            ));
+            acc.add_group(&group_label, assist_id, label, range, |builder| {
+                let scope = builder.make_import_scope_mut(scope.clone());
+                insert_use_as_alias(
+                    &scope,
+                    mod_path_to_ast(&import_path, edition),
+                    &ctx.config.insert_use,
+                    edition,
                 );
-            }
-            _ => {}
+            });
+        };
+
+        if let ImportCandidate::TraitAssocItem(name) | ImportCandidate::TraitMethod(name) =
+            import_assets.import_candidate()
+        {
+            if let hir::ItemInNs::Types(hir::ModuleDef::Trait(trait_to_import)) =
+                import.item_to_import
+                && trait_to_import.prefer_underscore_import(ctx.db())
+            {
+                // Flip the order of the suggestions and show a preference for `as _` in the name.
+                add_underscore_import(acc, name, format!("Import `{import_name}`"));
+                add_normal_import(acc, format!("Import `{import_name}` without `as _`"));
+            } else {
+                add_normal_import(acc, format!("Import `{import_name}`"));
+                add_underscore_import(acc, name, format!("Import `{import_name} as _`"));
+            }
+        } else {
+            add_normal_import(acc, format!("Import `{import_name}`"));
         }
     }
     Some(())
@@ -1957,4 +1961,72 @@ fn main() {
         "#,
         );
     }
+
+    #[test]
+    fn prefer_underscore_import() {
+        check_assist_by_label(
+            auto_import,
+            r#"
+mod foo {
+    #[rust_analyzer::prefer_underscore_import]
+    pub trait Ext {
+        fn bar(&self) {}
+    }
+    impl<T> Ext for T {}
+}
+
+fn baz() {
+    1.b$0ar();
+}
+        "#,
+            r#"
+use foo::Ext as _;
+
+mod foo {
+    #[rust_analyzer::prefer_underscore_import]
+    pub trait Ext {
+        fn bar(&self) {}
+    }
+    impl<T> Ext for T {}
+}
+
+fn baz() {
+    1.bar();
+}
+        "#,
+            "Import `foo::Ext`",
+        );
+        check_assist_by_label(
+            auto_import,
+            r#"
+mod foo {
+    #[rust_analyzer::prefer_underscore_import]
+    pub trait Ext {
+        fn bar(&self) {}
+    }
+    impl<T> Ext for T {}
+}
+
+fn baz() {
+    1.b$0ar();
+}
+        "#,
+            r#"
+use foo::Ext;
+
+mod foo {
+    #[rust_analyzer::prefer_underscore_import]
+    pub trait Ext {
+        fn bar(&self) {}
+    }
+    impl<T> Ext for T {}
+}
+
+fn baz() {
+    1.bar();
+}
+        "#,
+            "Import `foo::Ext` without `as _`",
+        );
+    }
 }

crates/ide-completion/src/item.rs

@@ -84,7 +84,15 @@ pub struct CompletionItem {
     pub ref_match: Option<(CompletionItemRefMode, TextSize)>,
 
     /// The import data to add to completion's edits.
-    pub import_to_add: SmallVec<[String; 1]>,
+    pub import_to_add: SmallVec<[CompletionItemImport; 1]>,
+}
+
+#[derive(Clone, UpmapFromRaFixture)]
+pub struct CompletionItemImport {
+    /// The path to import.
+    pub path: String,
+    /// Whether to import `as _`.
+    pub as_underscore: bool,
 }
 
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -585,7 +593,18 @@ impl Builder {
         let import_to_add = self
             .imports_to_add
             .into_iter()
-            .map(|import| import.import_path.display(db, self.edition).to_string())
+            .map(|import| {
+                let path = import.import_path.display(db, self.edition).to_string();
+                let as_underscore =
+                    if let hir::ItemInNs::Types(hir::ModuleDef::Trait(trait_to_import)) =
+                        import.item_to_import
+                    {
+                        trait_to_import.prefer_underscore_import(db)
+                    } else {
+                        false
+                    };
+                CompletionItemImport { path, as_underscore }
+            })
             .collect();
 
         CompletionItem {

crates/ide-completion/src/lib.rs

@@ -36,8 +36,8 @@ use crate::{
 pub use crate::{
     config::{AutoImportExclusionType, CallableSnippets, CompletionConfig},
     item::{
-        CompletionItem, CompletionItemKind, CompletionItemRefMode, CompletionRelevance,
-        CompletionRelevancePostfixMatch, CompletionRelevanceReturnType,
+        CompletionItem, CompletionItemImport, CompletionItemKind, CompletionItemRefMode,
+        CompletionRelevance, CompletionRelevancePostfixMatch, CompletionRelevanceReturnType,
         CompletionRelevanceTypeMatch,
     },
     snippet::{Snippet, SnippetScope},
@@ -280,7 +280,7 @@ pub fn resolve_completion_edits(
     db: &RootDatabase,
     config: &CompletionConfig<'_>,
     FilePosition { file_id, offset }: FilePosition,
-    imports: impl IntoIterator<Item = String>,
+    imports: impl IntoIterator<Item = CompletionItemImport>,
 ) -> Option<Vec<TextEdit>> {
     let _p = tracing::info_span!("resolve_completion_edits").entered();
     let sema = hir::Semantics::new(db);
@@ -299,12 +299,18 @@ pub fn resolve_completion_edits(
     let new_ast = scope.clone_for_update();
     let mut import_insert = TextEdit::builder();
 
-    imports.into_iter().for_each(|full_import_path| {
-        insert_use::insert_use(
-            &new_ast,
-            make::path_from_text_with_edition(&full_import_path, current_edition),
-            &config.insert_use,
-        );
+    imports.into_iter().for_each(|import| {
+        let full_path = make::path_from_text_with_edition(&import.path, current_edition);
+        if import.as_underscore {
+            insert_use::insert_use_as_alias(
+                &new_ast,
+                full_path,
+                &config.insert_use,
+                current_edition,
+            );
+        } else {
+            insert_use::insert_use(&new_ast, full_path, &config.insert_use);
+        }
     });
 
     diff(scope.as_syntax_node(), new_ast.as_syntax_node()).into_text_edit(&mut import_insert);

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

@@ -2057,3 +2057,38 @@ fn main() {
 "#,
     );
 }
+
+#[test]
+fn prefer_underscore_import() {
+    check_edit(
+        "bar",
+        r#"
+mod foo {
+    #[rust_analyzer::prefer_underscore_import]
+    pub trait Ext {
+        fn bar(&self) {}
+    }
+    impl<T> Ext for T {}
+}
+
+fn baz() {
+    1.bar$0
+}
+    "#,
+        r#"
+use foo::Ext as _;
+
+mod foo {
+    #[rust_analyzer::prefer_underscore_import]
+    pub trait Ext {
+        fn bar(&self) {}
+    }
+    impl<T> Ext for T {}
+}
+
+fn baz() {
+    1.bar();$0
+}
+    "#,
+    );
+}

crates/ide/src/lib.rs

@@ -127,7 +127,8 @@ pub use ide_assists::{
 };
 pub use ide_completion::{
     CallableSnippets, CompletionConfig, CompletionFieldsToResolve, CompletionItem,
-    CompletionItemKind, CompletionItemRefMode, CompletionRelevance, Snippet, SnippetScope,
+    CompletionItemImport, CompletionItemKind, CompletionItemRefMode, CompletionRelevance, Snippet,
+    SnippetScope,
 };
 pub use ide_db::{
     FileId, FilePosition, FileRange, RootDatabase, Severity, SymbolKind,
@@ -769,7 +770,7 @@ impl Analysis {
         &self,
         config: &CompletionConfig<'_>,
         position: FilePosition,
-        imports: impl IntoIterator<Item = String> + std::panic::UnwindSafe,
+        imports: impl IntoIterator<Item = CompletionItemImport> + std::panic::UnwindSafe,
     ) -> Cancellable<Vec<TextEdit>> {
         Ok(self
             .with_db(|db| ide_completion::resolve_completion_edits(db, config, position, imports))?

crates/rust-analyzer/src/handlers/request.rs

@@ -7,10 +7,10 @@ use anyhow::Context;
 
 use base64::{Engine, prelude::BASE64_STANDARD};
 use ide::{
-    AssistKind, AssistResolveStrategy, Cancellable, CompletionFieldsToResolve, FilePosition,
-    FileRange, FileStructureConfig, FindAllRefsConfig, HoverAction, HoverGotoTypeData,
-    InlayFieldsToResolve, Query, RangeInfo, Runnable, RunnableKind, SingleResolve, SourceChange,
-    TextEdit,
+    AssistKind, AssistResolveStrategy, Cancellable, CompletionFieldsToResolve,
+    CompletionItemImport, FilePosition, FileRange, FileStructureConfig, FindAllRefsConfig,
+    HoverAction, HoverGotoTypeData, InlayFieldsToResolve, Query, RangeInfo, Runnable, RunnableKind,
+    SingleResolve, SourceChange, TextEdit,
 };
 use ide_db::{FxHashMap, SymbolKind};
 use itertools::Itertools;
@@ -1233,7 +1233,10 @@ pub(crate) fn handle_completion_resolve(
             .resolve_completion_edits(
                 &forced_resolve_completions_config,
                 position,
-                resolve_data.imports.into_iter().map(|import| import.full_import_path),
+                resolve_data.imports.into_iter().map(|import| CompletionItemImport {
+                    path: import.full_import_path,
+                    as_underscore: import.as_underscore,
+                }),
             )?
             .into_iter()
             .flat_map(|edit| edit.into_iter().map(|indel| to_proto::text_edit(&line_index, indel)))

crates/rust-analyzer/src/lsp.rs

@@ -3,7 +3,7 @@
 use core::fmt;
 
 use hir::Mutability;
-use ide::{CompletionItem, CompletionItemRefMode, CompletionRelevance};
+use ide::{CompletionItem, CompletionItemImport, CompletionItemRefMode, CompletionRelevance};
 use tenthash::TentHash;
 
 pub mod ext;
@@ -136,8 +136,10 @@ pub(crate) fn completion_item_hash(item: &CompletionItem, is_ref_completion: boo
 
     hasher.update(item.import_to_add.len().to_ne_bytes());
     for import_path in &item.import_to_add {
-        hasher.update(import_path.len().to_ne_bytes());
-        hasher.update(import_path);
+        let CompletionItemImport { path, as_underscore } = import_path;
+        hasher.update(path.len().to_ne_bytes());
+        hasher.update(path);
+        hasher.update([u8::from(*as_underscore)]);
     }
 
     hasher.finalize()