1 files changed, 139 insertions, 8 deletions

diff --git a/src/tools/rust-analyzer/crates/ide-db/src/source_change.rs b/src/tools/rust-analyzer/crates/ide-db/src/source_change.rs
index 8132c73ef..8e338061d 100644
--- a/src/tools/rust-analyzer/crates/ide-db/src/source_change.rs
+++ b/src/tools/rust-analyzer/crates/ide-db/src/source_change.rs
@@ -3,16 +3,18 @@
 //!
 //! It can be viewed as a dual for `Change`.
 
-use std::{collections::hash_map::Entry, iter};
+use std::{collections::hash_map::Entry, iter, mem};
 
 use base_db::{AnchoredPathBuf, FileId};
-use rustc_hash::FxHashMap;
-use stdx::never;
-use text_edit::TextEdit;
+use stdx::{hash::NoHashHashMap, never};
+use syntax::{algo, AstNode, SyntaxNode, SyntaxNodePtr, TextRange, TextSize};
+use text_edit::{TextEdit, TextEditBuilder};
+
+use crate::SnippetCap;
 
 #[derive(Default, Debug, Clone)]
 pub struct SourceChange {
-    pub source_file_edits: FxHashMap<FileId, TextEdit>,
+    pub source_file_edits: NoHashHashMap<FileId, TextEdit>,
     pub file_system_edits: Vec<FileSystemEdit>,
     pub is_snippet: bool,
 }
@@ -21,7 +23,7 @@ impl SourceChange {
     /// Creates a new SourceChange with the given label
     /// from the edits.
     pub fn from_edits(
-        source_file_edits: FxHashMap<FileId, TextEdit>,
+        source_file_edits: NoHashHashMap<FileId, TextEdit>,
         file_system_edits: Vec<FileSystemEdit>,
     ) -> Self {
         SourceChange { source_file_edits, file_system_edits, is_snippet: false }
@@ -75,12 +77,141 @@ impl Extend<FileSystemEdit> for SourceChange {
     }
 }
 
-impl From<FxHashMap<FileId, TextEdit>> for SourceChange {
-    fn from(source_file_edits: FxHashMap<FileId, TextEdit>) -> SourceChange {
+impl From<NoHashHashMap<FileId, TextEdit>> for SourceChange {
+    fn from(source_file_edits: NoHashHashMap<FileId, TextEdit>) -> SourceChange {
         SourceChange { source_file_edits, file_system_edits: Vec::new(), is_snippet: false }
     }
 }
 
+pub struct SourceChangeBuilder {
+    pub edit: TextEditBuilder,
+    pub file_id: FileId,
+    pub source_change: SourceChange,
+    pub trigger_signature_help: bool,
+
+    /// Maps the original, immutable `SyntaxNode` to a `clone_for_update` twin.
+    pub mutated_tree: Option<TreeMutator>,
+}
+
+pub struct TreeMutator {
+    immutable: SyntaxNode,
+    mutable_clone: SyntaxNode,
+}
+
+impl TreeMutator {
+    pub fn new(immutable: &SyntaxNode) -> TreeMutator {
+        let immutable = immutable.ancestors().last().unwrap();
+        let mutable_clone = immutable.clone_for_update();
+        TreeMutator { immutable, mutable_clone }
+    }
+
+    pub fn make_mut<N: AstNode>(&self, node: &N) -> N {
+        N::cast(self.make_syntax_mut(node.syntax())).unwrap()
+    }
+
+    pub fn make_syntax_mut(&self, node: &SyntaxNode) -> SyntaxNode {
+        let ptr = SyntaxNodePtr::new(node);
+        ptr.to_node(&self.mutable_clone)
+    }
+}
+
+impl SourceChangeBuilder {
+    pub fn new(file_id: FileId) -> SourceChangeBuilder {
+        SourceChangeBuilder {
+            edit: TextEdit::builder(),
+            file_id,
+            source_change: SourceChange::default(),
+            trigger_signature_help: false,
+            mutated_tree: None,
+        }
+    }
+
+    pub fn edit_file(&mut self, file_id: FileId) {
+        self.commit();
+        self.file_id = file_id;
+    }
+
+    fn commit(&mut self) {
+        if let Some(tm) = self.mutated_tree.take() {
+            algo::diff(&tm.immutable, &tm.mutable_clone).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 fn make_mut<N: AstNode>(&mut self, node: N) -> N {
+        self.mutated_tree.get_or_insert_with(|| TreeMutator::new(node.syntax())).make_mut(&node)
+    }
+    /// 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 fn make_syntax_mut(&mut self, node: SyntaxNode) -> SyntaxNode {
+        self.mutated_tree.get_or_insert_with(|| TreeMutator::new(&node)).make_syntax_mut(&node)
+    }
+
+    /// Remove specified `range` of text.
+    pub fn delete(&mut self, range: TextRange) {
+        self.edit.delete(range)
+    }
+    /// Append specified `text` at the given `offset`
+    pub fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
+        self.edit.insert(offset, text.into())
+    }
+    /// Append specified `snippet` at the given `offset`
+    pub fn insert_snippet(
+        &mut self,
+        _cap: SnippetCap,
+        offset: TextSize,
+        snippet: impl Into<String>,
+    ) {
+        self.source_change.is_snippet = true;
+        self.insert(offset, snippet);
+    }
+    /// Replaces specified `range` of text with a given string.
+    pub fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
+        self.edit.replace(range, replace_with.into())
+    }
+    /// Replaces specified `range` of text with a given `snippet`.
+    pub fn replace_snippet(
+        &mut self,
+        _cap: SnippetCap,
+        range: TextRange,
+        snippet: impl Into<String>,
+    ) {
+        self.source_change.is_snippet = true;
+        self.replace(range, snippet);
+    }
+    pub fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
+        algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
+    }
+    pub fn create_file(&mut self, dst: AnchoredPathBuf, content: impl Into<String>) {
+        let file_system_edit = FileSystemEdit::CreateFile { dst, initial_contents: content.into() };
+        self.source_change.push_file_system_edit(file_system_edit);
+    }
+    pub fn move_file(&mut self, src: FileId, dst: AnchoredPathBuf) {
+        let file_system_edit = FileSystemEdit::MoveFile { src, dst };
+        self.source_change.push_file_system_edit(file_system_edit);
+    }
+    pub fn trigger_signature_help(&mut self) {
+        self.trigger_signature_help = true;
+    }
+
+    pub fn finish(mut self) -> SourceChange {
+        self.commit();
+        mem::take(&mut self.source_change)
+    }
+}
+
 #[derive(Debug, Clone)]
 pub enum FileSystemEdit {
     CreateFile { dst: AnchoredPathBuf, initial_contents: String },