Merge #8165

8165: Tweak assits API to fit mutable syntax trees r=matklad a=matklad

bors r+
🤖

Co-authored-by: Aleksey Kladov <[email protected]>

3 files changed, 53 insertions, 23 deletions

diff --git a/crates/ide_assists/src/assist_context.rs b/crates/ide_assists/src/assist_context.rs
index bba6c08e0..1482d37f8 100644
--- a/crates/ide_assists/src/assist_context.rs
+++ b/crates/ide_assists/src/assist_context.rs
@@ -14,8 +14,8 @@ use ide_db::{
 };
 use syntax::{
     algo::{self, find_node_at_offset, SyntaxRewriter},
-    AstNode, AstToken, SourceFile, SyntaxElement, SyntaxKind, SyntaxToken, TextRange, TextSize,
-    TokenAtOffset,
+    AstNode, AstToken, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxNodePtr,
+    SyntaxToken, TextRange, TextSize, TokenAtOffset,
 };
 use text_edit::{TextEdit, TextEditBuilder};
 
@@ -180,11 +180,19 @@ pub(crate) struct AssistBuilder {
     edit: TextEditBuilder,
     file_id: FileId,
     source_change: SourceChange,
+
+    /// Maps the original, immutable `SyntaxNode` to a `clone_for_update` twin.
+    mutated_tree: Option<(SyntaxNode, SyntaxNode)>,
 }
 
 impl AssistBuilder {
     pub(crate) fn new(file_id: FileId) -> AssistBuilder {
-        AssistBuilder { edit: TextEdit::builder(), file_id, source_change: SourceChange::default() }
+        AssistBuilder {
+            edit: TextEdit::builder(),
+            file_id,
+            source_change: SourceChange::default(),
+            mutated_tree: None,
+        }
     }
 
     pub(crate) fn edit_file(&mut self, file_id: FileId) {
@@ -193,12 +201,42 @@ impl AssistBuilder {
     }
 
     fn commit(&mut self) {
+        if let Some((old, new)) = self.mutated_tree.take() {
+            algo::diff(&old, &new).into_text_edit(&mut self.edit)
+        }
+
         let edit = mem::take(&mut self.edit).finish();
         if !edit.is_empty() {
             self.source_change.insert_source_edit(self.file_id, edit);
         }
     }
 
+    pub(crate) fn make_ast_mut<N: AstNode>(&mut self, node: N) -> N {
+        N::cast(self.make_mut(node.syntax().clone())).unwrap()
+    }
+    /// Returns a copy of the `node`, suitable for mutation.
+    ///
+    /// Syntax trees in rust-analyzer are typically immutable, and mutating
+    /// operations panic at runtime. However, it is possible to make a copy of
+    /// the tree and mutate the copy freely. Mutation is based on interior
+    /// mutability, and different nodes in the same tree see the same mutations.
+    ///
+    /// The typical pattern for an assist is to find specific nodes in the read
+    /// phase, and then get their mutable couterparts using `make_mut` in the
+    /// mutable state.
+    pub(crate) fn make_mut(&mut self, node: SyntaxNode) -> SyntaxNode {
+        let root = &self
+            .mutated_tree
+            .get_or_insert_with(|| {
+                let immutable = node.ancestors().last().unwrap();
+                let mutable = immutable.clone_for_update();
+                (immutable, mutable)
+            })
+            .1;
+        let ptr = SyntaxNodePtr::new(&&node);
+        ptr.to_node(root)
+    }
+
     /// Remove specified `range` of text.
     pub(crate) fn delete(&mut self, range: TextRange) {
         self.edit.delete(range)
diff --git a/crates/ide_assists/src/handlers/merge_imports.rs b/crates/ide_assists/src/handlers/merge_imports.rs
index cfc472a32..8e0794218 100644
--- a/crates/ide_assists/src/handlers/merge_imports.rs
+++ b/crates/ide_assists/src/handlers/merge_imports.rs
@@ -21,12 +21,6 @@ use crate::{
 // ```
 pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let tree: ast::UseTree = ctx.find_node_at_offset()?;
-    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;
@@ -53,22 +47,20 @@ pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext) -> Option<()
         target,
         |builder| {
             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());
+                let to_replace = builder.make_ast_mut(to_replace);
+                let to_remove = builder.make_ast_mut(to_remove);
+
+                ted::replace(to_replace.syntax(), replacement.syntax());
                 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());
+                let to_replace = builder.make_ast_mut(to_replace);
+                let to_remove = builder.make_ast_mut(to_remove);
+
+                ted::replace(to_replace.syntax(), replacement.syntax());
                 to_remove.remove()
             }
-
-            builder.replace(original_parent.text_range(), new_parent.to_string())
         },
     )
 }
diff --git a/crates/ide_assists/src/handlers/move_bounds.rs b/crates/ide_assists/src/handlers/move_bounds.rs
index b5dec8014..011a28d44 100644
--- a/crates/ide_assists/src/handlers/move_bounds.rs
+++ b/crates/ide_assists/src/handlers/move_bounds.rs
@@ -21,7 +21,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // }
 // ```
 pub(crate) fn move_bounds_to_where_clause(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let type_param_list = ctx.find_node_at_offset::<ast::GenericParamList>()?.clone_for_update();
+    let type_param_list = ctx.find_node_at_offset::<ast::GenericParamList>()?;
 
     let mut type_params = type_param_list.type_params();
     if type_params.all(|p| p.type_bound_list().is_none()) {
@@ -29,7 +29,6 @@ pub(crate) fn move_bounds_to_where_clause(acc: &mut Assists, ctx: &AssistContext
     }
 
     let parent = type_param_list.syntax().parent()?;
-    let original_parent_range = parent.text_range();
 
     let target = type_param_list.syntax().text_range();
     acc.add(
@@ -37,6 +36,9 @@ pub(crate) fn move_bounds_to_where_clause(acc: &mut Assists, ctx: &AssistContext
         "Move to where clause",
         target,
         |edit| {
+            let type_param_list = edit.make_ast_mut(type_param_list);
+            let parent = edit.make_mut(parent);
+
             let where_clause: ast::WhereClause = match_ast! {
                 match parent {
                     ast::Fn(it) => it.get_or_create_where_clause(),
@@ -56,8 +58,6 @@ pub(crate) fn move_bounds_to_where_clause(acc: &mut Assists, ctx: &AssistContext
                     tbl.remove()
                 }
             }
-
-            edit.replace(original_parent_range, parent.to_string())
         },
     )
 }