From f42e02aaff84d7b613ab9057036e93010c23b93f Mon Sep 17 00:00:00 2001 From: Chayim Refael Friedman Date: Wed, 4 Mar 2026 05:54:57 +0200 Subject: [PATCH] Allow crate authors to declare that their trait prefers to be imported `as _` For example for extension traits. Provide an attribute for that. It'll affect flyimport and the autoimport quickfix, as explained in the code. --- crates/hir-def/src/attrs.rs | 5 + crates/hir/src/lib.rs | 13 ++ .../ide-assists/src/handlers/auto_import.rs | 144 +++++++++++++----- crates/ide-completion/src/item.rs | 23 ++- crates/ide-completion/src/lib.rs | 24 +-- crates/ide-completion/src/tests/flyimport.rs | 35 +++++ crates/ide/src/lib.rs | 5 +- crates/rust-analyzer/src/handlers/request.rs | 13 +- crates/rust-analyzer/src/lsp.rs | 8 +- crates/rust-analyzer/src/lsp/ext.rs | 1 + crates/rust-analyzer/src/lsp/to_proto.rs | 5 +- docs/book/src/contributing/lsp-extensions.md | 2 +- 12 files changed, 219 insertions(+), 59 deletions(-) diff --git a/crates/hir-def/src/attrs.rs b/crates/hir-def/src/attrs.rs index 92d492889386..5cf5a9b6be84 100644 --- a/crates/hir-def/src/attrs.rs +++ b/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; } } diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs index 27516ed80b27..cd3391b77132 100644 --- a/crates/hir/src/lib.rs +++ b/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 { diff --git a/crates/ide-assists/src/handlers/auto_import.rs b/crates/ide-assists/src/handlers/auto_import.rs index adeb191719fb..f9d618790c6e 100644 --- a/crates/ide-assists/src/handlers/auto_import.rs +++ b/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 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 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 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 Ext for T {} +} + +fn baz() { + 1.bar(); +} + "#, + "Import `foo::Ext` without `as _`", + ); + } } diff --git a/crates/ide-completion/src/item.rs b/crates/ide-completion/src/item.rs index 62211a808ca6..6abf4f632aa8 100644 --- a/crates/ide-completion/src/item.rs +++ b/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 { diff --git a/crates/ide-completion/src/lib.rs b/crates/ide-completion/src/lib.rs index 3867e65ae57e..3df511a5ad0f 100644 --- a/crates/ide-completion/src/lib.rs +++ b/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, + imports: impl IntoIterator, ) -> Option> { 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); diff --git a/crates/ide-completion/src/tests/flyimport.rs b/crates/ide-completion/src/tests/flyimport.rs index 5391e6c9ce6e..60ae077d0142 100644 --- a/crates/ide-completion/src/tests/flyimport.rs +++ b/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 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 Ext for T {} +} + +fn baz() { + 1.bar();$0 +} + "#, + ); +} diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs index 8a6ef6418063..442dc6cfbfb2 100644 --- a/crates/ide/src/lib.rs +++ b/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 + std::panic::UnwindSafe, + imports: impl IntoIterator + std::panic::UnwindSafe, ) -> Cancellable> { Ok(self .with_db(|db| ide_completion::resolve_completion_edits(db, config, position, imports))? diff --git a/crates/rust-analyzer/src/handlers/request.rs b/crates/rust-analyzer/src/handlers/request.rs index 4193bd25b210..0c1c067ffa91 100644 --- a/crates/rust-analyzer/src/handlers/request.rs +++ b/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))) diff --git a/crates/rust-analyzer/src/lsp.rs b/crates/rust-analyzer/src/lsp.rs index c7a5a95e66bb..a6a35dadd991 100644 --- a/crates/rust-analyzer/src/lsp.rs +++ b/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() diff --git a/crates/rust-analyzer/src/lsp/ext.rs b/crates/rust-analyzer/src/lsp/ext.rs index e6493eefef17..5d0d9209de2f 100644 --- a/crates/rust-analyzer/src/lsp/ext.rs +++ b/crates/rust-analyzer/src/lsp/ext.rs @@ -858,6 +858,7 @@ pub struct InlayHintResolveData { #[derive(Debug, Serialize, Deserialize)] pub struct CompletionImport { pub full_import_path: String, + pub as_underscore: bool, } #[derive(Debug, Deserialize, Default)] diff --git a/crates/rust-analyzer/src/lsp/to_proto.rs b/crates/rust-analyzer/src/lsp/to_proto.rs index 5fa95252e7cb..d857f23b6703 100644 --- a/crates/rust-analyzer/src/lsp/to_proto.rs +++ b/crates/rust-analyzer/src/lsp/to_proto.rs @@ -413,7 +413,10 @@ fn completion_item( item.import_to_add .clone() .into_iter() - .map(|import_path| lsp_ext::CompletionImport { full_import_path: import_path }) + .map(|import| lsp_ext::CompletionImport { + full_import_path: import.path, + as_underscore: import.as_underscore, + }) .collect() } else { Vec::new() diff --git a/docs/book/src/contributing/lsp-extensions.md b/docs/book/src/contributing/lsp-extensions.md index 22c1784ac293..28fb979b01b3 100644 --- a/docs/book/src/contributing/lsp-extensions.md +++ b/docs/book/src/contributing/lsp-extensions.md @@ -1,5 +1,5 @@