Merge #6465

6465: Support multiple file edits in AssistBuilder r=matklad a=Veykril

Fixes #6459

Co-authored-by: Lukas Wirth <[email protected]>

2 files changed, 34 insertions, 14 deletions

diff --git a/crates/assists/src/assist_context.rs b/crates/assists/src/assist_context.rs
index 51a160f40..a17e592b0 100644
--- a/crates/assists/src/assist_context.rs
+++ b/crates/assists/src/assist_context.rs
@@ -208,7 +208,7 @@ pub(crate) struct AssistBuilder {
     edit: TextEditBuilder,
     file_id: FileId,
     is_snippet: bool,
-    change: SourceChange,
+    source_file_edits: Vec<SourceFileEdit>,
 }
 
 impl AssistBuilder {
@@ -217,20 +217,27 @@ impl AssistBuilder {
             edit: TextEdit::builder(),
             file_id,
             is_snippet: false,
-            change: SourceChange::default(),
+            source_file_edits: Vec::default(),
         }
     }
 
     pub(crate) fn edit_file(&mut self, file_id: FileId) {
+        self.commit();
         self.file_id = file_id;
     }
 
     fn commit(&mut self) {
         let edit = mem::take(&mut self.edit).finish();
         if !edit.is_empty() {
-            let new_edit = SourceFileEdit { file_id: self.file_id, edit };
-            assert!(!self.change.source_file_edits.iter().any(|it| it.file_id == new_edit.file_id));
-            self.change.source_file_edits.push(new_edit);
+            match self.source_file_edits.binary_search_by_key(&self.file_id, |edit| edit.file_id) {
+                Ok(idx) => self.source_file_edits[idx]
+                    .edit
+                    .union(edit)
+                    .expect("overlapping edits for same file"),
+                Err(idx) => self
+                    .source_file_edits
+                    .insert(idx, SourceFileEdit { file_id: self.file_id, edit }),
+            }
         }
     }
 
@@ -277,10 +284,10 @@ impl AssistBuilder {
 
     fn finish(mut self) -> SourceChange {
         self.commit();
-        let mut change = mem::take(&mut self.change);
-        if self.is_snippet {
-            change.is_snippet = true;
+        SourceChange {
+            source_file_edits: mem::take(&mut self.source_file_edits),
+            file_system_edits: Default::default(),
+            is_snippet: self.is_snippet,
         }
-        change
     }
 }
diff --git a/editors/code/src/snippets.ts b/editors/code/src/snippets.ts
index 258b49982..fee736e7d 100644
--- a/editors/code/src/snippets.ts
+++ b/editors/code/src/snippets.ts
@@ -3,16 +3,29 @@ import * as vscode from 'vscode';
 import { assert } from './util';
 
 export async function applySnippetWorkspaceEdit(edit: vscode.WorkspaceEdit) {
-    assert(edit.entries().length === 1, `bad ws edit: ${JSON.stringify(edit)}`);
-    const [uri, edits] = edit.entries()[0];
+    if (edit.entries().length === 1) {
+        const [uri, edits] = edit.entries()[0];
+        const editor = await editorFromUri(uri);
+        if (editor) await applySnippetTextEdits(editor, edits);
+        return;
+    }
+    for (const [uri, edits] of edit.entries()) {
+        const editor = await editorFromUri(uri);
+        if (editor) await editor.edit((builder) => {
+            for (const indel of edits) {
+                assert(!parseSnippet(indel.newText), `bad ws edit: snippet received with multiple edits: ${JSON.stringify(edit)}`);
+                builder.replace(indel.range, indel.newText);
+            }
+        });
+    }
+}
 
+async function editorFromUri(uri: vscode.Uri): Promise<vscode.TextEditor | undefined> {
     if (vscode.window.activeTextEditor?.document.uri !== uri) {
         // `vscode.window.visibleTextEditors` only contains editors whose contents are being displayed
         await vscode.window.showTextDocument(uri, {});
     }
-    const editor = vscode.window.visibleTextEditors.find((it) => it.document.uri.toString() === uri.toString());
-    if (!editor) return;
-    await applySnippetTextEdits(editor, edits);
+    return vscode.window.visibleTextEditors.find((it) => it.document.uri.toString() === uri.toString());
 }
 
 export async function applySnippetTextEdits(editor: vscode.TextEditor, edits: vscode.TextEdit[]) {