feat: import trait with alias

2 files changed, 77 insertions, 12 deletions

diff --git a/crates/ide-assists/src/handlers/auto_import.rs b/crates/ide-assists/src/handlers/auto_import.rs
index 7acf2ea0a0d..a2efa4d10bd 100644
--- a/crates/ide-assists/src/handlers/auto_import.rs
+++ b/crates/ide-assists/src/handlers/auto_import.rs
@@ -5,7 +5,7 @@ use ide_db::{
     helpers::mod_path_to_ast,
     imports::{
         import_assets::{ImportAssets, ImportCandidate, LocatedImport},
-        insert_use::{insert_use, ImportScope},
+        insert_use::{insert_use, insert_use_as_alias, ImportScope},
     },
 };
 use syntax::{ast, AstNode, NodeOrToken, SyntaxElement};
@@ -129,10 +129,12 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
     for import in proposed_imports {
         let import_path = import.import_path;
 
+        let (assist_id, import_name) =
+            (AssistId("auto_import", AssistKind::QuickFix), import_path.display(ctx.db()));
         acc.add_group(
             &group_label,
-            AssistId("auto_import", AssistKind::QuickFix),
-            format!("Import `{}`", import_path.display(ctx.db())),
+            assist_id,
+            format!("Import `{}`", import_name),
             range,
             |builder| {
                 let scope = match scope.clone() {
@@ -143,6 +145,38 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
                 insert_use(&scope, mod_path_to_ast(&import_path), &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 `{} as _`", import_name),
+                    range,
+                    |builder| {
+                        let scope = match scope.clone() {
+                            ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
+                            ImportScope::Module(it) => ImportScope::Module(builder.make_mut(it)),
+                            ImportScope::Block(it) => ImportScope::Block(builder.make_mut(it)),
+                        };
+                        insert_use_as_alias(
+                            &scope,
+                            mod_path_to_ast(&import_path),
+                            &ctx.config.insert_use,
+                        );
+                    },
+                );
+            }
+            _ => {}
+        }
     }
     Some(())
 }
@@ -723,7 +757,7 @@ fn main() {
             }
             ",
             r"
-            use test_mod::TestTrait;
+            use test_mod::TestTrait as _;
 
             mod test_mod {
                 pub trait TestTrait {
@@ -794,7 +828,7 @@ fn main() {
             }
             ",
             r"
-            use test_mod::TestTrait;
+            use test_mod::TestTrait as _;
 
             mod test_mod {
                 pub trait TestTrait {
@@ -866,7 +900,7 @@ fn main() {
             }
             ",
             r"
-            use test_mod::TestTrait;
+            use test_mod::TestTrait as _;
 
             mod test_mod {
                 pub trait TestTrait {
@@ -908,7 +942,7 @@ fn main() {
             }
             ",
             r"
-            use dep::test_mod::TestTrait;
+            use dep::test_mod::TestTrait as _;
 
             fn main() {
                 let test_struct = dep::test_mod::TestStruct {};
@@ -939,7 +973,7 @@ fn main() {
             }
             ",
             r"
-            use dep::test_mod::TestTrait;
+            use dep::test_mod::TestTrait as _;
 
             fn main() {
                 dep::test_mod::TestStruct::test_function
@@ -969,7 +1003,7 @@ fn main() {
             }
             ",
             r"
-            use dep::test_mod::TestTrait;
+            use dep::test_mod::TestTrait as _;
 
             fn main() {
                 dep::test_mod::TestStruct::CONST
diff --git a/crates/ide-db/src/imports/insert_use.rs b/crates/ide-db/src/imports/insert_use.rs
index 9be1d366349..6caa6e3511c 100644
--- a/crates/ide-db/src/imports/insert_use.rs
+++ b/crates/ide-db/src/imports/insert_use.rs
@@ -9,7 +9,7 @@ use syntax::{
     algo,
     ast::{
         self, edit_in_place::Removable, make, AstNode, HasAttrs, HasModuleItem, HasVisibility,
-        PathSegmentKind,
+        PathSegmentKind, Rename, UseTree,
     },
     ted, Direction, NodeOrToken, SyntaxKind, SyntaxNode,
 };
@@ -157,6 +157,33 @@ impl ImportScope {
 
 /// Insert an import path into the given file/node. A `merge` value of none indicates that no import merging is allowed to occur.
 pub fn insert_use(scope: &ImportScope, path: ast::Path, cfg: &InsertUseConfig) {
+    insert_use_with_alias_option(scope, path, cfg, None);
+}
+
+pub fn insert_use_as_alias(scope: &ImportScope, path: ast::Path, cfg: &InsertUseConfig) {
+    let text: &str = "use foo as _";
+    let parse = syntax::SourceFile::parse(text);
+    let node = match parse.tree().syntax().descendants().find_map(UseTree::cast) {
+        Some(it) => it,
+        None => {
+            panic!(
+                "Failed to make ast node `{}` from text {}",
+                std::any::type_name::<Rename>(),
+                text
+            )
+        }
+    };
+    let alias = node.rename();
+
+    insert_use_with_alias_option(scope, path, cfg, alias);
+}
+
+fn insert_use_with_alias_option(
+    scope: &ImportScope,
+    path: ast::Path,
+    cfg: &InsertUseConfig,
+    alias: Option<ast::Rename>,
+) {
     let _p = profile::span("insert_use");
     let mut mb = match cfg.granularity {
         ImportGranularity::Crate => Some(MergeBehavior::Crate),
@@ -175,8 +202,12 @@ pub fn insert_use(scope: &ImportScope, path: ast::Path, cfg: &InsertUseConfig) {
         };
     }
 
-    let use_item =
-        make::use_(None, make::use_tree(path.clone(), None, None, false)).clone_for_update();
+    let use_item = if alias.is_some() {
+        make::use_(None, make::use_tree(path.clone(), None, alias, false)).clone_for_update()
+    } else {
+        make::use_(None, make::use_tree(path.clone(), None, None, false)).clone_for_update()
+    };
+
     // merge into existing imports if possible
     if let Some(mb) = mb {
         let filter = |it: &_| !(cfg.skip_glob_imports && ast::Use::is_simple_glob(it));