2 files changed, 42 insertions, 37 deletions

diff --git a/crates/ide_assists/src/handlers/merge_imports.rs b/crates/ide_assists/src/handlers/merge_imports.rs
index 7bd7e1e36..cfc472a32 100644
--- a/crates/ide_assists/src/handlers/merge_imports.rs
+++ b/crates/ide_assists/src/handlers/merge_imports.rs
@@ -1,8 +1,5 @@
 use ide_db::helpers::insert_use::{try_merge_imports, try_merge_trees, MergeBehavior};
-use syntax::{
-    algo::{neighbor, SyntaxRewriter},
-    ast, AstNode,
-};
+use syntax::{algo::neighbor, ast, ted, AstNode};
 
 use crate::{
     assist_context::{AssistContext, Assists},
@@ -24,33 +21,29 @@ use crate::{
 // ```
 pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let tree: ast::UseTree = ctx.find_node_at_offset()?;
-    let mut rewriter = SyntaxRewriter::default();
+    let original_parent = tree.syntax().ancestors().last()?;
+
+    let tree = tree.clone_for_update();
+    let new_parent = tree.syntax().ancestors().last()?;
+
     let mut offset = ctx.offset();
 
+    let mut imports = None;
+    let mut uses = None;
     if let Some(use_item) = tree.syntax().parent().and_then(ast::Use::cast) {
-        let (merged, to_delete) =
+        let (merged, to_remove) =
             next_prev().filter_map(|dir| neighbor(&use_item, dir)).find_map(|use_item2| {
                 try_merge_imports(&use_item, &use_item2, MergeBehavior::Full).zip(Some(use_item2))
             })?;
 
-        rewriter.replace_ast(&use_item, &merged);
-        rewriter += to_delete.remove();
-
-        if to_delete.syntax().text_range().end() < offset {
-            offset -= to_delete.syntax().text_range().len();
-        }
+        imports = Some((use_item, merged, to_remove));
     } else {
-        let (merged, to_delete) =
+        let (merged, to_remove) =
             next_prev().filter_map(|dir| neighbor(&tree, dir)).find_map(|use_tree| {
                 try_merge_trees(&tree, &use_tree, MergeBehavior::Full).zip(Some(use_tree))
             })?;
 
-        rewriter.replace_ast(&tree, &merged);
-        rewriter += to_delete.remove();
-
-        if to_delete.syntax().text_range().end() < offset {
-            offset -= to_delete.syntax().text_range().len();
-        }
+        uses = Some((tree.clone(), merged, to_remove))
     };
 
     let target = tree.syntax().text_range();
@@ -59,7 +52,23 @@ pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext) -> Option<()
         "Merge imports",
         target,
         |builder| {
-            builder.rewrite(rewriter);
+            if let Some((to_replace, replacement, to_remove)) = imports {
+                if to_remove.syntax().text_range().end() < offset {
+                    offset -= to_remove.syntax().text_range().len();
+                }
+                ted::replace(to_replace.syntax().clone(), replacement.syntax().clone());
+                to_remove.remove();
+            }
+
+            if let Some((to_replace, replacement, to_remove)) = uses {
+                if to_remove.syntax().text_range().end() < offset {
+                    offset -= to_remove.syntax().text_range().len();
+                }
+                ted::replace(to_replace.syntax().clone(), replacement.syntax().clone());
+                to_remove.remove()
+            }
+
+            builder.replace(original_parent.text_range(), new_parent.to_string())
         },
     )
 }
diff --git a/crates/ide_assists/src/handlers/unmerge_use.rs b/crates/ide_assists/src/handlers/unmerge_use.rs
index 616af7c2e..8d271e056 100644
--- a/crates/ide_assists/src/handlers/unmerge_use.rs
+++ b/crates/ide_assists/src/handlers/unmerge_use.rs
@@ -1,6 +1,6 @@
 use syntax::{
-    algo::SyntaxRewriter,
-    ast::{self, edit::AstNodeEdit, VisibilityOwner},
+    ast::{self, VisibilityOwner},
+    ted::{self, Position},
     AstNode, SyntaxKind,
 };
 
@@ -22,7 +22,7 @@ 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()?;
+    let tree: ast::UseTree = ctx.find_node_at_offset::<ast::UseTree>()?.clone_for_update();
 
     let tree_list = tree.syntax().parent().and_then(ast::UseTreeList::cast)?;
     if tree_list.use_trees().count() < 2 {
@@ -33,6 +33,9 @@ pub(crate) fn unmerge_use(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     let use_: ast::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()?;
+
     let target = tree.syntax().text_range();
     acc.add(
         AssistId("unmerge_use", AssistKind::RefactorRewrite),
@@ -47,20 +50,13 @@ pub(crate) fn unmerge_use(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
                     tree.rename(),
                     tree.star_token().is_some(),
                 ),
-            );
-
-            let mut rewriter = SyntaxRewriter::default();
-            rewriter += tree.remove();
-            rewriter.insert_after(use_.syntax(), &ast::make::tokens::single_newline());
-            if let ident_level @ 1..=usize::MAX = use_.indent_level().0 as usize {
-                rewriter.insert_after(
-                    use_.syntax(),
-                    &ast::make::tokens::whitespace(&" ".repeat(4 * ident_level)),
-                );
-            }
-            rewriter.insert_after(use_.syntax(), new_use.syntax());
-
-            builder.rewrite(rewriter);
+            )
+            .clone_for_update();
+
+            tree.remove();
+            ted::insert(Position::after(use_.syntax()), new_use.syntax());
+
+            builder.replace(old_parent_range, new_parent.to_string());
         },
     )
 }