fix: migrate unmerge_use to syntax editor

Also ensures that attributes on the use item are applied to the new use
item when unmerging.

Signed-off-by: Prajwal S N <[email protected]>

Author: snprajwal

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

@@ -1,7 +1,13 @@
 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, Indent},
+        make,
+        syntax_factory::SyntaxFactory,
+    },
+    syntax_editor::{Element, Position, Removable},
 };
 
 use crate::{
@@ -22,20 +28,17 @@ use crate::{
 // 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();
+    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);
         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}`"),
@@ -44,16 +47,28 @@ 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_(
+        let make = SyntaxFactory::new();
+        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());
+        });
+        // Indent the new use item
+        new_use.indent(IndentLevel::from_node(use_.syntax()));
 
-        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 leading whitespace and the new use item after the current use item
+        editor.insert_all(
+            Position::after(use_.syntax()),
+            vec![make.whitespace("\n").syntax_element(), new_use.syntax().syntax_element()],
+        );
+        editor.add_mappings(make.finish_with_mappings());
+        builder.add_file_edits(ctx.file_id(), editor);
     })
 }
 
@@ -230,4 +245,19 @@ pub use std::fmt::Display;
 use std::process;",
         );
     }
+
+    #[test]
+    fn unmerge_use_item_with_attributes() {
+        check_assist(
+            unmerge_use,
+            r"
+#[allow(deprecated)]
+use foo::{bar, baz$0};",
+            r"
+#[allow(deprecated)]
+use foo::{bar};
+#[allow(deprecated)]
+use foo::baz;",
+        );
+    }
 }

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();
 

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)]

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::new();
+
+        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;