Auto merge of #16475 - DropDemBits:native-vscode-snippet-text-edit, r=Veykril

feat: Support multiple tab stops for completions in VSCode

Uses the native VSCode support for `SnippetTextEdit`s. Fixes #13229 and fixes #8531.

https://github.com/rust-lang/rust-analyzer/assets/13354275/a2d2c033-bb30-4f34-92ca-bf3f4f744cdc

This is done in a slightly hacky way, as `vscode-languageclient` can't convert RA's `SnippetTextEdit`s into vscode `SnippetTextEdit`s and will appear to use a [different format](295aaa393f/types/src/main.ts (L1501-L1516)) in the future.

---

~~Marked as draft since as-is, this will cause completions to double-indent any multi-line code generated.~~
**Update:** This also fixes up edits so that any multi-line code won't be double-indented.

editors/code/src/commands.ts

@@ -4,7 +4,11 @@ import * as ra from "./lsp_ext";
 import * as path from "path";
 
 import type { Ctx, Cmd, CtxInit } from "./ctx";
-import { applySnippetWorkspaceEdit, applySnippetTextEdits } from "./snippets";
+import {
+    applySnippetWorkspaceEdit,
+    applySnippetTextEdits,
+    type SnippetTextDocumentEdit,
+} from "./snippets";
 import { spawnSync } from "child_process";
 import { type RunnableQuickPick, selectRunnable, createTask, createArgs } from "./run";
 import { AstInspector } from "./ast_inspector";
@@ -1006,7 +1010,6 @@ export function resolveCodeAction(ctx: CtxInit): Cmd {
             return;
         }
         const itemEdit = item.edit;
-        const edit = await client.protocol2CodeConverter.asWorkspaceEdit(itemEdit);
         // filter out all text edits and recreate the WorkspaceEdit without them so we can apply
         // snippet edits on our own
         const lcFileSystemEdit = {
@@ -1017,16 +1020,71 @@ export function resolveCodeAction(ctx: CtxInit): Cmd {
             lcFileSystemEdit,
         );
         await vscode.workspace.applyEdit(fileSystemEdit);
-        await applySnippetWorkspaceEdit(edit);
+
+        // replace all text edits so that we can convert snippet text edits into `vscode.SnippetTextEdit`s
+        // FIXME: this is a workaround until vscode-languageclient supports doing the SnippeTextEdit conversion itself
+        // also need to carry the snippetTextDocumentEdits separately, since we can't retrieve them again using WorkspaceEdit.entries
+        const [workspaceTextEdit, snippetTextDocumentEdits] = asWorkspaceSnippetEdit(ctx, itemEdit);
+        await applySnippetWorkspaceEdit(workspaceTextEdit, snippetTextDocumentEdits);
         if (item.command != null) {
             await vscode.commands.executeCommand(item.command.command, item.command.arguments);
         }
     };
 }
 
+function asWorkspaceSnippetEdit(
+    ctx: CtxInit,
+    item: lc.WorkspaceEdit,
+): [vscode.WorkspaceEdit, SnippetTextDocumentEdit[]] {
+    const client = ctx.client;
+
+    // partially borrowed from https://github.com/microsoft/vscode-languageserver-node/blob/295aaa393fda8ecce110c38880a00466b9320e63/client/src/common/protocolConverter.ts#L1060-L1101
+    const result = new vscode.WorkspaceEdit();
+
+    if (item.documentChanges) {
+        const snippetTextDocumentEdits: SnippetTextDocumentEdit[] = [];
+
+        for (const change of item.documentChanges) {
+            if (lc.TextDocumentEdit.is(change)) {
+                const uri = client.protocol2CodeConverter.asUri(change.textDocument.uri);
+                const snippetTextEdits: (vscode.TextEdit | vscode.SnippetTextEdit)[] = [];
+
+                for (const edit of change.edits) {
+                    if (
+                        "insertTextFormat" in edit &&
+                        edit.insertTextFormat === lc.InsertTextFormat.Snippet
+                    ) {
+                        // is a snippet text edit
+                        snippetTextEdits.push(
+                            new vscode.SnippetTextEdit(
+                                client.protocol2CodeConverter.asRange(edit.range),
+                                new vscode.SnippetString(edit.newText),
+                            ),
+                        );
+                    } else {
+                        // always as a text document edit
+                        snippetTextEdits.push(
+                            vscode.TextEdit.replace(
+                                client.protocol2CodeConverter.asRange(edit.range),
+                                edit.newText,
+                            ),
+                        );
+                    }
+                }
+
+                snippetTextDocumentEdits.push([uri, snippetTextEdits]);
+            }
+        }
+        return [result, snippetTextDocumentEdits];
+    } else {
+        // we don't handle WorkspaceEdit.changes since it's not relevant for code actions
+        return [result, []];
+    }
+}
+
 export function applySnippetWorkspaceEditCommand(_ctx: CtxInit): Cmd {
     return async (edit: vscode.WorkspaceEdit) => {
-        await applySnippetWorkspaceEdit(edit);
+        await applySnippetWorkspaceEdit(edit, edit.entries());
     };
 }
 

editors/code/src/snippets.ts

@@ -3,20 +3,28 @@ import * as vscode from "vscode";
 import { assert } from "./util";
 import { unwrapUndefinable } from "./undefinable";
 
-export async function applySnippetWorkspaceEdit(edit: vscode.WorkspaceEdit) {
+export type SnippetTextDocumentEdit = [vscode.Uri, (vscode.TextEdit | vscode.SnippetTextEdit)[]];
-    if (edit.entries().length === 1) {
+
-        const [uri, edits] = unwrapUndefinable(edit.entries()[0]);
+export async function applySnippetWorkspaceEdit(
+    edit: vscode.WorkspaceEdit,
+    editEntries: SnippetTextDocumentEdit[],
+) {
+    if (editEntries.length === 1) {
+        const [uri, edits] = unwrapUndefinable(editEntries[0]);
         const editor = await editorFromUri(uri);
-        if (editor) await applySnippetTextEdits(editor, edits);
+        if (editor) {
+            edit.set(uri, removeLeadingWhitespace(editor, edits));
+            await vscode.workspace.applyEdit(edit);
+        }
         return;
     }
-    for (const [uri, edits] of edit.entries()) {
+    for (const [uri, edits] of editEntries) {
         const editor = await editorFromUri(uri);
         if (editor) {
             await editor.edit((builder) => {
                 for (const indel of edits) {
                     assert(
-                        !parseSnippet(indel.newText),
+                        !(indel instanceof vscode.SnippetTextEdit),
                         `bad ws edit: snippet received with multiple edits: ${JSON.stringify(
                             edit,
                         )}`,
@@ -39,53 +47,97 @@ async function editorFromUri(uri: vscode.Uri): Promise<vscode.TextEditor | undef
 }
 
 export async function applySnippetTextEdits(editor: vscode.TextEditor, edits: vscode.TextEdit[]) {
-    const selections: vscode.Selection[] = [];
+    const edit = new vscode.WorkspaceEdit();
-    let lineDelta = 0;
+    const snippetEdits = toSnippetTextEdits(edits);
-    await editor.edit((builder) => {
+    edit.set(editor.document.uri, removeLeadingWhitespace(editor, snippetEdits));
-        for (const indel of edits) {
+    await vscode.workspace.applyEdit(edit);
-            const parsed = parseSnippet(indel.newText);
-            if (parsed) {
-                const [newText, [placeholderStart, placeholderLength]] = parsed;
-                const prefix = newText.substr(0, placeholderStart);
-                const lastNewline = prefix.lastIndexOf("\n");
-
-                const startLine = indel.range.start.line + lineDelta + countLines(prefix);
-                const startColumn =
-                    lastNewline === -1
-                        ? indel.range.start.character + placeholderStart
-                        : prefix.length - lastNewline - 1;
-                const endColumn = startColumn + placeholderLength;
-                selections.push(
-                    new vscode.Selection(
-                        new vscode.Position(startLine, startColumn),
-                        new vscode.Position(startLine, endColumn),
-                    ),
-                );
-                builder.replace(indel.range, newText);
-            } else {
-                builder.replace(indel.range, indel.newText);
 }
-            lineDelta +=
+
-                countLines(indel.newText) - (indel.range.end.line - indel.range.start.line);
+function hasSnippet(snip: string): boolean {
+    const m = snip.match(/\$\d+|\{\d+:[^}]*\}/);
+    return m != null;
+}
+
+function toSnippetTextEdits(
+    edits: vscode.TextEdit[],
+): (vscode.TextEdit | vscode.SnippetTextEdit)[] {
+    return edits.map((textEdit) => {
+        // Note: text edits without any snippets are returned as-is instead of
+        // being wrapped in a SnippetTextEdit, as otherwise it would be
+        // treated as if it had a tab stop at the end.
+        if (hasSnippet(textEdit.newText)) {
+            return new vscode.SnippetTextEdit(
+                textEdit.range,
+                new vscode.SnippetString(textEdit.newText),
+            );
+        } else {
+            return textEdit;
         }
     });
-    if (selections.length > 0) editor.selections = selections;
-    if (selections.length === 1) {
-        const selection = unwrapUndefinable(selections[0]);
-        editor.revealRange(selection, vscode.TextEditorRevealType.InCenterIfOutsideViewport);
-    }
 }
 
-function parseSnippet(snip: string): [string, [number, number]] | undefined {
+/**
-    const m = snip.match(/\$(0|\{0:([^}]*)\})/);
+ * Removes the leading whitespace from snippet edits, so as to not double up
-    if (!m) return undefined;
+ * on indentation.
-    const placeholder = m[2] ?? "";
+ *
-    if (m.index == null) return undefined;
+ * Snippet edits by default adjust any multi-line snippets to match the
-    const range: [number, number] = [m.index, placeholder.length];
+ * indentation of the line to insert at. Unfortunately, we (the server) also
-    const insert = snip.replace(m[0], placeholder);
+ * include the required indentation to match what we line insert at, so we end
-    return [insert, range];
+ * up doubling up the indentation. Since there isn't any way to tell vscode to
+ * not fixup indentation for us, we instead opt to remove the indentation and
+ * then let vscode add it back in.
+ *
+ * This assumes that the source snippet text edits have the required
+ * indentation, but that's okay as even without this workaround and the problem
+ * to workaround, those snippet edits would already be inserting at the wrong
+ * indentation.
+ */
+function removeLeadingWhitespace(
+    editor: vscode.TextEditor,
+    edits: (vscode.TextEdit | vscode.SnippetTextEdit)[],
+) {
+    return edits.map((edit) => {
+        if (edit instanceof vscode.SnippetTextEdit) {
+            const snippetEdit: vscode.SnippetTextEdit = edit;
+            const firstLineEnd = snippetEdit.snippet.value.indexOf("\n");
+
+            if (firstLineEnd !== -1) {
+                // Is a multi-line snippet, remove the indentation which
+                // would be added back in by vscode.
+                const startLine = editor.document.lineAt(snippetEdit.range.start.line);
+                const leadingWhitespace = getLeadingWhitespace(
+                    startLine.text,
+                    0,
+                    startLine.firstNonWhitespaceCharacterIndex,
+                );
+
+                const [firstLine, rest] = splitAt(snippetEdit.snippet.value, firstLineEnd + 1);
+                const unindentedLines = rest
+                    .split("\n")
+                    .map((line) => line.replace(leadingWhitespace, ""))
+                    .join("\n");
+
+                snippetEdit.snippet.value = firstLine + unindentedLines;
             }
 
-function countLines(text: string): number {
+            return snippetEdit;
-    return (text.match(/\n/g) || []).length;
+        } else {
+            return edit;
+        }
+    });
+}
+
+// based on https://github.com/microsoft/vscode/blob/main/src/vs/base/common/strings.ts#L284
+function getLeadingWhitespace(str: string, start: number = 0, end: number = str.length): string {
+    for (let i = start; i < end; i++) {
+        const chCode = str.charCodeAt(i);
+        if (chCode !== " ".charCodeAt(0) && chCode !== " ".charCodeAt(0)) {
+            return str.substring(start, i);
+        }
+    }
+    return str.substring(start, end);
+}
+
+function splitAt(str: string, index: number): [string, string] {
+    return [str.substring(0, index), str.substring(index)];
 }