Merge pull request #19469 from snprajwal/merge-imports

refactor: `merge_imports` and `unmerge_imports` to editor

7 files changed, 175 insertions, 101 deletions

diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/merge_imports.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/merge_imports.rs
index b7f7cb9cb01..6bf7f584914 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/merge_imports.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/merge_imports.rs
@@ -1,14 +1,14 @@
 use either::Either;
 use ide_db::imports::{
     insert_use::{ImportGranularity, InsertUseConfig},
-    merge_imports::{MergeBehavior, try_merge_imports, try_merge_trees, try_normalize_use_tree},
+    merge_imports::{MergeBehavior, try_merge_imports, try_merge_trees},
 };
-use itertools::Itertools;
 use syntax::{
     AstNode, SyntaxElement, SyntaxNode,
     algo::neighbor,
-    ast::{self, edit_in_place::Removable},
-    match_ast, ted,
+    ast::{self, syntax_factory::SyntaxFactory},
+    match_ast,
+    syntax_editor::Removable,
 };
 
 use crate::{
@@ -69,49 +69,32 @@ pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext<'_>) -> Optio
         (selection_range, edits?)
     };
 
+    let parent_node = match ctx.covering_element() {
+        SyntaxElement::Node(n) => n,
+        SyntaxElement::Token(t) => t.parent()?,
+    };
+
     acc.add(AssistId::refactor_rewrite("merge_imports"), "Merge imports", target, |builder| {
-        let edits_mut: Vec<Edit> = edits
-            .into_iter()
-            .map(|it| match it {
-                Remove(Either::Left(it)) => Remove(Either::Left(builder.make_mut(it))),
-                Remove(Either::Right(it)) => Remove(Either::Right(builder.make_mut(it))),
-                Replace(old, new) => Replace(builder.make_syntax_mut(old), new),
-            })
-            .collect();
-        for edit in edits_mut {
+        let make = SyntaxFactory::with_mappings();
+        let mut editor = builder.make_editor(&parent_node);
+
+        for edit in edits {
             match edit {
-                Remove(it) => it.as_ref().either(Removable::remove, Removable::remove),
-                Replace(old, new) => {
-                    ted::replace(old, &new);
-
-                    // If there's a selection and we're replacing a use tree in a tree list,
-                    // normalize the parent use tree if it only contains the merged subtree.
-                    if !ctx.has_empty_selection() {
-                        let normalized_use_tree = ast::UseTree::cast(new)
-                            .as_ref()
-                            .and_then(ast::UseTree::parent_use_tree_list)
-                            .and_then(|use_tree_list| {
-                                if use_tree_list.use_trees().collect_tuple::<(_,)>().is_some() {
-                                    Some(use_tree_list.parent_use_tree())
-                                } else {
-                                    None
-                                }
-                            })
-                            .and_then(|target_tree| {
-                                try_normalize_use_tree(
-                                    &target_tree,
-                                    ctx.config.insert_use.granularity.into(),
-                                )
-                                .map(|top_use_tree_flat| (target_tree, top_use_tree_flat))
-                            });
-                        if let Some((old_tree, new_tree)) = normalized_use_tree {
-                            cov_mark::hit!(replace_parent_with_normalized_use_tree);
-                            ted::replace(old_tree.syntax(), new_tree.syntax());
-                        }
+                Remove(it) => {
+                    let node = it.as_ref();
+                    if let Some(left) = node.left() {
+                        left.remove(&mut editor);
+                    } else if let Some(right) = node.right() {
+                        right.remove(&mut editor);
                     }
                 }
+                Replace(old, new) => {
+                    editor.replace(old, &new);
+                }
             }
         }
+        editor.add_mappings(make.finish_with_mappings());
+        builder.add_file_edits(ctx.vfs_file_id(), editor);
     })
 }
 
@@ -723,11 +706,10 @@ use std::{
         );
 
         cov_mark::check!(merge_with_selected_use_tree_neighbors);
-        cov_mark::check!(replace_parent_with_normalized_use_tree);
         check_assist(
             merge_imports,
             r"use std::$0{fmt::Display, fmt::Debug}$0;",
-            r"use std::fmt::{Debug, Display};",
+            r"use std::{fmt::{Debug, Display}};",
         );
     }
 
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unmerge_use.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unmerge_imports.rs
index 805a7344494..c066f41ca47 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unmerge_use.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unmerge_imports.rs
@@ -1,7 +1,10 @@
 use syntax::{
     AstNode, SyntaxKind,
-    ast::{self, HasVisibility, edit_in_place::Removable, make},
-    ted::{self, Position},
+    ast::{
+        self, HasAttrs, HasVisibility, edit::IndentLevel, edit_in_place::AttrsOwnerEdit, make,
+        syntax_factory::SyntaxFactory,
+    },
+    syntax_editor::{Element, Position, Removable},
 };
 
 use crate::{
@@ -9,9 +12,9 @@ use crate::{
     assist_context::{AssistContext, Assists},
 };
 
-// Assist: unmerge_use
+// Assist: unmerge_imports
 //
-// Extracts single use item from use list.
+// Extracts a use item from a use list into a standalone use list.
 //
 // ```
 // use std::fmt::{Debug, Display$0};
@@ -21,21 +24,18 @@ use crate::{
 // use std::fmt::{Debug};
 // use std::fmt::Display;
 // ```
-pub(crate) fn unmerge_use(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    let tree: ast::UseTree = ctx.find_node_at_offset::<ast::UseTree>()?.clone_for_update();
+pub(crate) fn unmerge_imports(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let tree = ctx.find_node_at_offset::<ast::UseTree>()?;
 
     let tree_list = tree.syntax().parent().and_then(ast::UseTreeList::cast)?;
     if tree_list.use_trees().count() < 2 {
-        cov_mark::hit!(skip_single_use_item);
+        cov_mark::hit!(skip_single_import);
         return None;
     }
 
-    let use_: ast::Use = tree_list.syntax().ancestors().find_map(ast::Use::cast)?;
+    let use_ = tree_list.syntax().ancestors().find_map(ast::Use::cast)?;
     let path = resolve_full_path(&tree)?;
 
-    let old_parent_range = use_.syntax().parent()?.text_range();
-    let new_parent = use_.syntax().parent()?;
-
     // If possible, explain what is going to be done.
     let label = match tree.path().and_then(|path| path.first_segment()) {
         Some(name) => format!("Unmerge use of `{name}`"),
@@ -43,17 +43,31 @@ pub(crate) fn unmerge_use(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
     };
 
     let target = tree.syntax().text_range();
-    acc.add(AssistId::refactor_rewrite("unmerge_use"), label, target, |builder| {
-        let new_use = make::use_(
+    acc.add(AssistId::refactor_rewrite("unmerge_imports"), label, target, |builder| {
+        let make = SyntaxFactory::with_mappings();
+        let new_use = make.use_(
             use_.visibility(),
-            make::use_tree(path, tree.use_tree_list(), tree.rename(), tree.star_token().is_some()),
-        )
-        .clone_for_update();
-
-        tree.remove();
-        ted::insert(Position::after(use_.syntax()), new_use.syntax());
+            make.use_tree(path, tree.use_tree_list(), tree.rename(), tree.star_token().is_some()),
+        );
+        // Add any attributes that are present on the use tree
+        use_.attrs().for_each(|attr| {
+            new_use.add_attr(attr.clone_for_update());
+        });
 
-        builder.replace(old_parent_range, new_parent.to_string());
+        let mut editor = builder.make_editor(use_.syntax());
+        // Remove the use tree from the current use item
+        tree.remove(&mut editor);
+        // Insert a newline and indentation, followed by the new use item
+        editor.insert_all(
+            Position::after(use_.syntax()),
+            vec![
+                make.whitespace(&format!("\n{}", IndentLevel::from_node(use_.syntax())))
+                    .syntax_element(),
+                new_use.syntax().syntax_element(),
+            ],
+        );
+        editor.add_mappings(make.finish_with_mappings());
+        builder.add_file_edits(ctx.vfs_file_id(), editor);
     })
 }
 
@@ -80,22 +94,22 @@ mod tests {
     use super::*;
 
     #[test]
-    fn skip_single_use_item() {
-        cov_mark::check!(skip_single_use_item);
+    fn skip_single_import() {
+        cov_mark::check!(skip_single_import);
         check_assist_not_applicable(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::Debug$0;
 ",
         );
         check_assist_not_applicable(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::{Debug$0};
 ",
         );
         check_assist_not_applicable(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::Debug as Dbg$0;
 ",
@@ -105,7 +119,7 @@ use std::fmt::Debug as Dbg$0;
     #[test]
     fn skip_single_glob_import() {
         check_assist_not_applicable(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::*$0;
 ",
@@ -113,9 +127,9 @@ use std::fmt::*$0;
     }
 
     #[test]
-    fn unmerge_use_item() {
+    fn unmerge_import() {
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::{Debug, Display$0};
 ",
@@ -126,7 +140,7 @@ use std::fmt::Display;
         );
 
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::{Debug, format$0, Display};
 ",
@@ -140,7 +154,7 @@ use std::fmt::format;
     #[test]
     fn unmerge_glob_import() {
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::{*$0, Display};
 ",
@@ -152,9 +166,9 @@ use std::fmt::*;
     }
 
     #[test]
-    fn unmerge_renamed_use_item() {
+    fn unmerge_renamed_import() {
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 use std::fmt::{Debug, Display as Disp$0};
 ",
@@ -166,9 +180,9 @@ use std::fmt::Display as Disp;
     }
 
     #[test]
-    fn unmerge_indented_use_item() {
+    fn unmerge_indented_import() {
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 mod format {
     use std::fmt::{Debug, Display$0 as Disp, format};
@@ -184,9 +198,9 @@ mod format {
     }
 
     #[test]
-    fn unmerge_nested_use_item() {
+    fn unmerge_nested_import() {
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 use foo::bar::{baz::{qux$0, foobar}, barbaz};
 ",
@@ -196,7 +210,7 @@ use foo::bar::baz::qux;
 ",
         );
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 use foo::bar::{baz$0::{qux, foobar}, barbaz};
 ",
@@ -208,9 +222,9 @@ use foo::bar::baz::{qux, foobar};
     }
 
     #[test]
-    fn unmerge_use_item_with_visibility() {
+    fn unmerge_import_with_visibility() {
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"
 pub use std::fmt::{Debug, Display$0};
 ",
@@ -222,12 +236,27 @@ pub use std::fmt::Display;
     }
 
     #[test]
-    fn unmerge_use_item_on_self() {
+    fn unmerge_import_on_self() {
         check_assist(
-            unmerge_use,
+            unmerge_imports,
             r"use std::process::{Command, self$0};",
             r"use std::process::{Command};
 use std::process;",
         );
     }
+
+    #[test]
+    fn unmerge_import_with_attributes() {
+        check_assist(
+            unmerge_imports,
+            r"
+#[allow(deprecated)]
+use foo::{bar, baz$0};",
+            r"
+#[allow(deprecated)]
+use foo::{bar};
+#[allow(deprecated)]
+use foo::baz;",
+        );
+    }
 }
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
index a157483a449..627ed37b04e 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
@@ -222,8 +222,8 @@ mod handlers {
     mod toggle_async_sugar;
     mod toggle_ignore;
     mod toggle_macro_delimiter;
+    mod unmerge_imports;
     mod unmerge_match_arm;
-    mod unmerge_use;
     mod unnecessary_async;
     mod unqualify_method_call;
     mod unwrap_block;
@@ -363,7 +363,7 @@ mod handlers {
             toggle_ignore::toggle_ignore,
             toggle_macro_delimiter::toggle_macro_delimiter,
             unmerge_match_arm::unmerge_match_arm,
-            unmerge_use::unmerge_use,
+            unmerge_imports::unmerge_imports,
             unnecessary_async::unnecessary_async,
             unqualify_method_call::unqualify_method_call,
             unwrap_block::unwrap_block,
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs b/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
index 00a9d35c310..54d060cc790 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
@@ -3340,6 +3340,20 @@ sth!{ }
 }
 
 #[test]
+fn doctest_unmerge_imports() {
+    check_doc_test(
+        "unmerge_imports",
+        r#####"
+use std::fmt::{Debug, Display$0};
+"#####,
+        r#####"
+use std::fmt::{Debug};
+use std::fmt::Display;
+"#####,
+    )
+}
+
+#[test]
 fn doctest_unmerge_match_arm() {
     check_doc_test(
         "unmerge_match_arm",
@@ -3366,20 +3380,6 @@ fn handle(action: Action) {
 }
 
 #[test]
-fn doctest_unmerge_use() {
-    check_doc_test(
-        "unmerge_use",
-        r#####"
-use std::fmt::{Debug, Display$0};
-"#####,
-        r#####"
-use std::fmt::{Debug};
-use std::fmt::Display;
-"#####,
-    )
-}
-
-#[test]
 fn doctest_unnecessary_async() {
     check_doc_test(
         "unnecessary_async",
diff --git a/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs b/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs
index 1854000d3db..1894a3218f8 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs
@@ -107,6 +107,20 @@ impl SyntaxFactory {
         ast
     }
 
+    pub fn use_(&self, visibility: Option<ast::Visibility>, use_tree: ast::UseTree) -> ast::Use {
+        make::use_(visibility, use_tree).clone_for_update()
+    }
+
+    pub fn use_tree(
+        &self,
+        path: ast::Path,
+        use_tree_list: Option<ast::UseTreeList>,
+        alias: Option<ast::Rename>,
+        add_star: bool,
+    ) -> ast::UseTree {
+        make::use_tree(path, use_tree_list, alias, add_star).clone_for_update()
+    }
+
     pub fn path_unqualified(&self, segment: ast::PathSegment) -> ast::Path {
         let ast = make::path_unqualified(segment.clone()).clone_for_update();
 
diff --git a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs
index 58200189c46..31caf618be9 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor.rs
@@ -20,6 +20,7 @@ mod edit_algo;
 mod edits;
 mod mapping;
 
+pub use edits::Removable;
 pub use mapping::{SyntaxMapping, SyntaxMappingBuilder};
 
 #[derive(Debug)]
diff --git a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edits.rs b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edits.rs
index 350cb3e2544..d66ea8aa28c 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edits.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/syntax_editor/edits.rs
@@ -1,7 +1,8 @@
 //! Structural editing for ast using `SyntaxEditor`
 
 use crate::{
-    Direction, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, T,
+    AstToken, Direction, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, T,
+    algo::neighbor,
     ast::{
         self, AstNode, Fn, GenericParam, HasGenericParams, HasName, edit::IndentLevel, make,
         syntax_factory::SyntaxFactory,
@@ -143,6 +144,53 @@ fn normalize_ws_between_braces(editor: &mut SyntaxEditor, node: &SyntaxNode) ->
     Some(())
 }
 
+pub trait Removable: AstNode {
+    fn remove(&self, editor: &mut SyntaxEditor);
+}
+
+impl Removable for ast::Use {
+    fn remove(&self, editor: &mut SyntaxEditor) {
+        let make = SyntaxFactory::without_mappings();
+
+        let next_ws = self
+            .syntax()
+            .next_sibling_or_token()
+            .and_then(|it| it.into_token())
+            .and_then(ast::Whitespace::cast);
+        if let Some(next_ws) = next_ws {
+            let ws_text = next_ws.syntax().text();
+            if let Some(rest) = ws_text.strip_prefix('\n') {
+                if rest.is_empty() {
+                    editor.delete(next_ws.syntax());
+                } else {
+                    editor.replace(next_ws.syntax(), make.whitespace(rest));
+                }
+            }
+        }
+
+        editor.delete(self.syntax());
+    }
+}
+
+impl Removable for ast::UseTree {
+    fn remove(&self, editor: &mut SyntaxEditor) {
+        for dir in [Direction::Next, Direction::Prev] {
+            if let Some(next_use_tree) = neighbor(self, dir) {
+                let separators = self
+                    .syntax()
+                    .siblings_with_tokens(dir)
+                    .skip(1)
+                    .take_while(|it| it.as_node() != Some(next_use_tree.syntax()));
+                for sep in separators {
+                    editor.delete(sep);
+                }
+                break;
+            }
+        }
+        editor.delete(self.syntax());
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use parser::Edition;