internal: Defer rendering of structured snippets

This ensures that any assist using structured snippets won't
accidentally remove bits interpreted as snippet bits.

4 files changed, 161 insertions, 78 deletions

diff --git a/crates/ide-assists/src/tests.rs b/crates/ide-assists/src/tests.rs
index 00cea0e76c6..cc3e251a8f2 100644
--- a/crates/ide-assists/src/tests.rs
+++ b/crates/ide-assists/src/tests.rs
@@ -132,8 +132,13 @@ fn check_doc_test(assist_id: &str, before: &str, after: &str) {
             .filter(|it| !it.source_file_edits.is_empty() || !it.file_system_edits.is_empty())
             .expect("Assist did not contain any source changes");
         let mut actual = before;
-        if let Some(source_file_edit) = source_change.get_source_edit(file_id) {
+        if let Some((source_file_edit, snippet_edit)) =
+            source_change.get_source_and_snippet_edit(file_id)
+        {
             source_file_edit.apply(&mut actual);
+            if let Some(snippet_edit) = snippet_edit {
+                snippet_edit.apply(&mut actual);
+            }
         }
         actual
     };
@@ -191,9 +196,12 @@ fn check_with_config(
                 && source_change.file_system_edits.len() == 0;
 
             let mut buf = String::new();
-            for (file_id, (edit, _snippet_edit)) in source_change.source_file_edits {
+            for (file_id, (edit, snippet_edit)) in source_change.source_file_edits {
                 let mut text = db.file_text(file_id).as_ref().to_owned();
                 edit.apply(&mut text);
+                if let Some(snippet_edit) = snippet_edit {
+                    snippet_edit.apply(&mut text);
+                }
                 if !skip_header {
                     let sr = db.file_source_root(file_id);
                     let sr = db.source_root(sr);
diff --git a/crates/ide-db/src/source_change.rs b/crates/ide-db/src/source_change.rs
index 596f28e9816..3ff56ae9027 100644
--- a/crates/ide-db/src/source_change.rs
+++ b/crates/ide-db/src/source_change.rs
@@ -11,8 +11,7 @@ use itertools::Itertools;
 use nohash_hasher::IntMap;
 use stdx::never;
 use syntax::{
-    algo, ast, ted, AstNode, SyntaxElement, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextRange,
-    TextSize,
+    algo, AstNode, SyntaxElement, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextRange, TextSize,
 };
 use text_edit::{TextEdit, TextEditBuilder};
 
@@ -76,8 +75,11 @@ impl SourceChange {
         self.file_system_edits.push(edit);
     }
 
-    pub fn get_source_edit(&self, file_id: FileId) -> Option<&TextEdit> {
-        self.source_file_edits.get(&file_id).map(|(edit, _)| edit)
+    pub fn get_source_and_snippet_edit(
+        &self,
+        file_id: FileId,
+    ) -> Option<&(TextEdit, Option<SnippetEdit>)> {
+        self.source_file_edits.get(&file_id)
     }
 
     pub fn merge(mut self, other: SourceChange) -> SourceChange {
@@ -258,24 +260,19 @@ impl SourceChangeBuilder {
     }
 
     fn commit(&mut self) {
-        // Render snippets first so that they get bundled into the tree diff
-        if let Some(mut snippets) = self.snippet_builder.take() {
-            // Last snippet always has stop index 0
-            let last_stop = snippets.places.pop().unwrap();
-            last_stop.place(0);
-
-            for (index, stop) in snippets.places.into_iter().enumerate() {
-                stop.place(index + 1)
-            }
-        }
+        let snippet_edit = self.snippet_builder.take().map(|builder| {
+            SnippetEdit::new(
+                builder.places.into_iter().map(PlaceSnippet::finalize_position).collect_vec(),
+            )
+        });
 
         if let Some(tm) = self.mutated_tree.take() {
-            algo::diff(&tm.immutable, &tm.mutable_clone).into_text_edit(&mut self.edit)
+            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);
+        if !edit.is_empty() || snippet_edit.is_some() {
+            self.source_change.insert_source_and_snippet_edit(self.file_id, edit, snippet_edit);
         }
     }
 
@@ -429,57 +426,11 @@ enum PlaceSnippet {
 }
 
 impl PlaceSnippet {
-    /// Places the snippet before or over an element with the given tab stop index
-    fn place(self, order: usize) {
-        // ensure the target element is still attached
-        match &self {
-            PlaceSnippet::Before(element)
-            | PlaceSnippet::After(element)
-            | PlaceSnippet::Over(element) => {
-                // element should still be in the tree, but if it isn't
-                // then it's okay to just ignore this place
-                if stdx::never!(element.parent().is_none()) {
-                    return;
-                }
-            }
-        }
-
+    fn finalize_position(self) -> Snippet {
         match self {
-            PlaceSnippet::Before(element) => {
-                ted::insert_raw(ted::Position::before(&element), Self::make_tab_stop(order));
-            }
-            PlaceSnippet::After(element) => {
-                ted::insert_raw(ted::Position::after(&element), Self::make_tab_stop(order));
-            }
-            PlaceSnippet::Over(element) => {
-                let position = ted::Position::before(&element);
-                element.detach();
-
-                let snippet = ast::SourceFile::parse(&format!("${{{order}:_}}"))
-                    .syntax_node()
-                    .clone_for_update();
-
-                let placeholder =
-                    snippet.descendants().find_map(ast::UnderscoreExpr::cast).unwrap();
-                ted::replace(placeholder.syntax(), element);
-
-                ted::insert_raw(position, snippet);
-            }
+            PlaceSnippet::Before(it) => Snippet::Tabstop(it.text_range().start()),
+            PlaceSnippet::After(it) => Snippet::Tabstop(it.text_range().end()),
+            PlaceSnippet::Over(it) => Snippet::Placeholder(it.text_range()),
         }
     }
-
-    fn make_tab_stop(order: usize) -> SyntaxNode {
-        let stop = ast::SourceFile::parse(&format!("stop!(${order})"))
-            .syntax_node()
-            .descendants()
-            .find_map(ast::TokenTree::cast)
-            .unwrap()
-            .syntax()
-            .clone_for_update();
-
-        stop.first_token().unwrap().detach();
-        stop.last_token().unwrap().detach();
-
-        stop
-    }
 }
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 0ad4c6c47e6..bf77d55d58e 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -127,7 +127,7 @@ pub use ide_db::{
     label::Label,
     line_index::{LineCol, LineIndex},
     search::{ReferenceCategory, SearchScope},
-    source_change::{FileSystemEdit, SourceChange},
+    source_change::{FileSystemEdit, SnippetEdit, SourceChange},
     symbol_index::Query,
     RootDatabase, SymbolKind,
 };
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index 0022b7a8674..46ca7db2e16 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -10,8 +10,8 @@ use ide::{
     CompletionItemKind, CompletionRelevance, Documentation, FileId, FileRange, FileSystemEdit,
     Fold, FoldKind, Highlight, HlMod, HlOperator, HlPunct, HlRange, HlTag, Indel, InlayHint,
     InlayHintLabel, InlayHintLabelPart, InlayKind, Markup, NavigationTarget, ReferenceCategory,
-    RenameError, Runnable, Severity, SignatureHelp, SourceChange, StructureNodeKind, SymbolKind,
-    TextEdit, TextRange, TextSize,
+    RenameError, Runnable, Severity, SignatureHelp, SnippetEdit, SourceChange, StructureNodeKind,
+    SymbolKind, TextEdit, TextRange, TextSize,
 };
 use itertools::Itertools;
 use serde_json::to_value;
@@ -22,7 +22,7 @@ use crate::{
     config::{CallInfoConfig, Config},
     global_state::GlobalStateSnapshot,
     line_index::{LineEndings, LineIndex, PositionEncoding},
-    lsp_ext,
+    lsp_ext::{self, SnippetTextEdit},
     lsp_utils::invalid_params_error,
     semantic_tokens::{self, standard_fallback_type},
 };
@@ -884,16 +884,135 @@ fn outside_workspace_annotation_id() -> String {
     String::from("OutsideWorkspace")
 }
 
+fn merge_text_and_snippet_edit(
+    line_index: &LineIndex,
+    edit: TextEdit,
+    snippet_edit: Option<SnippetEdit>,
+) -> Vec<SnippetTextEdit> {
+    let Some(snippet_edit) = snippet_edit else {
+        return edit.into_iter().map(|it| snippet_text_edit(&line_index, false, it)).collect();
+    };
+
+    let mut edits: Vec<SnippetTextEdit> = vec![];
+    let mut snippets = snippet_edit.into_edit_ranges().into_iter().peekable();
+    let mut text_edits = edit.into_iter();
+
+    while let Some(current_indel) = text_edits.next() {
+        let new_range = {
+            let insert_len =
+                TextSize::try_from(current_indel.insert.len()).unwrap_or(TextSize::from(u32::MAX));
+            TextRange::at(current_indel.delete.start(), insert_len)
+        };
+
+        // insert any snippets before the text edit
+        let first_snippet_in_or_after_edit = loop {
+            let Some((snippet_index, snippet_range)) = snippets.peek() else { break None };
+
+            // check if we're entirely before the range
+            // only possible for tabstops
+            if snippet_range.end() < new_range.start()
+                && stdx::always!(
+                    snippet_range.is_empty(),
+                    "placeholder range is before any text edits"
+                )
+            {
+                let range = range(&line_index, *snippet_range);
+                let new_text = format!("${snippet_index}");
+
+                edits.push(SnippetTextEdit {
+                    range,
+                    new_text,
+                    insert_text_format: Some(lsp_types::InsertTextFormat::SNIPPET),
+                    annotation_id: None,
+                })
+            } else {
+                break Some((snippet_index, snippet_range));
+            }
+        };
+
+        if first_snippet_in_or_after_edit
+            .is_some_and(|(_, range)| new_range.intersect(*range).is_some())
+        {
+            // at least one snippet edit intersects this text edit,
+            // so gather all of the edits that intersect this text edit
+            let mut all_snippets = snippets
+                .take_while_ref(|(_, range)| new_range.intersect(*range).is_some())
+                .collect_vec();
+
+            // ensure all of the ranges are wholly contained inside of the new range
+            all_snippets.retain(|(_, range)| {
+                    stdx::always!(
+                        new_range.contains_range(*range),
+                        "found placeholder range {:?} which wasn't fully inside of text edit's new range {:?}", range, new_range
+                    )
+                });
+
+            let mut text_edit = text_edit(line_index, current_indel);
+
+            // escape out snippet text
+            stdx::replace(&mut text_edit.new_text, '\\', r"\\");
+            stdx::replace(&mut text_edit.new_text, '$', r"\$");
+
+            // ...and apply!
+            for (index, range) in all_snippets.iter().rev() {
+                let start = (range.start() - new_range.start()).into();
+                let end = (range.end() - new_range.start()).into();
+
+                if range.is_empty() {
+                    text_edit.new_text.insert_str(start, &format!("${index}"));
+                } else {
+                    text_edit.new_text.insert(end, '}');
+                    text_edit.new_text.insert_str(start, &format!("${{{index}"));
+                }
+            }
+
+            edits.push(SnippetTextEdit {
+                range: text_edit.range,
+                new_text: text_edit.new_text,
+                insert_text_format: Some(lsp_types::InsertTextFormat::SNIPPET),
+                annotation_id: None,
+            })
+        } else {
+            // snippet edit was beyond the current one
+            // since it wasn't consumed, it's available for the next pass
+            edits.push(snippet_text_edit(line_index, false, current_indel));
+        }
+    }
+
+    // insert any remaining edits
+    // either one of the two or both should've run out at this point,
+    // so it's either a tail of text edits or tabstops
+    edits.extend(text_edits.map(|indel| snippet_text_edit(line_index, false, indel)));
+    edits.extend(snippets.map(|(snippet_index, snippet_range)| {
+        stdx::always!(
+            snippet_range.is_empty(),
+            "found placeholder snippet {:?} without a text edit",
+            snippet_range
+        );
+
+        let range = range(&line_index, snippet_range);
+        let new_text = format!("${snippet_index}");
+
+        SnippetTextEdit {
+            range,
+            new_text,
+            insert_text_format: Some(lsp_types::InsertTextFormat::SNIPPET),
+            annotation_id: None,
+        }
+    }));
+
+    edits
+}
+
 pub(crate) fn snippet_text_document_edit(
     snap: &GlobalStateSnapshot,
-    is_snippet: bool,
     file_id: FileId,
     edit: TextEdit,
+    snippet_edit: Option<SnippetEdit>,
 ) -> Cancellable<lsp_ext::SnippetTextDocumentEdit> {
     let text_document = optional_versioned_text_document_identifier(snap, file_id);
     let line_index = snap.file_line_index(file_id)?;
-    let mut edits: Vec<_> =
-        edit.into_iter().map(|it| snippet_text_edit(&line_index, is_snippet, it)).collect();
+    let mut edits = merge_text_and_snippet_edit(&line_index, edit, snippet_edit);
 
     if snap.analysis.is_library_file(file_id)? && snap.config.change_annotation_support() {
         for edit in &mut edits {
@@ -973,8 +1092,13 @@ pub(crate) fn snippet_workspace_edit(
         let ops = snippet_text_document_ops(snap, op)?;
         document_changes.extend_from_slice(&ops);
     }
-    for (file_id, (edit, _snippet_edit)) in source_change.source_file_edits {
-        let edit = snippet_text_document_edit(snap, source_change.is_snippet, file_id, edit)?;
+    for (file_id, (edit, snippet_edit)) in source_change.source_file_edits {
+        let edit = snippet_text_document_edit(
+            snap,
+            file_id,
+            edit,
+            snippet_edit.filter(|_| source_change.is_snippet),
+        )?;
         document_changes.push(lsp_ext::SnippetDocumentChangeOperation::Edit(edit));
     }
     let mut workspace_edit = lsp_ext::SnippetWorkspaceEdit {