4 files changed, 53 insertions, 31 deletions

diff --git a/editors/code/src/commands/index.ts b/editors/code/src/commands/index.ts
index 0937b495c..e5ed77e32 100644
--- a/editors/code/src/commands/index.ts
+++ b/editors/code/src/commands/index.ts
@@ -53,36 +53,57 @@ export function selectAndApplySourceChange(ctx: Ctx): Cmd {
     };
 }
 
-export function applySnippetWorkspaceEdit(_ctx: Ctx): Cmd {
+export function applySnippetWorkspaceEditCommand(_ctx: Ctx): Cmd {
     return async (edit: vscode.WorkspaceEdit) => {
-        assert(edit.entries().length === 1, `bad ws edit: ${JSON.stringify(edit)}`);
-        const [uri, edits] = edit.entries()[0];
+        await applySnippetWorkspaceEdit(edit);
+    };
+}
+
+export async function applySnippetWorkspaceEdit(edit: vscode.WorkspaceEdit) {
+    assert(edit.entries().length === 1, `bad ws edit: ${JSON.stringify(edit)}`);
+    const [uri, edits] = edit.entries()[0];
 
-        const editor = vscode.window.visibleTextEditors.find((it) => it.document.uri.toString() === uri.toString());
-        if (!editor) return;
+    const editor = vscode.window.visibleTextEditors.find((it) => it.document.uri.toString() === uri.toString());
+    if (!editor) return;
 
-        let editWithSnippet: vscode.TextEdit | undefined = undefined;
-        let lineDelta = 0;
-        await editor.edit((builder) => {
-            for (const indel of edits) {
-                const isSnippet = indel.newText.indexOf('$0') !== -1 || indel.newText.indexOf('${') !== -1;
-                if (isSnippet) {
-                    editWithSnippet = indel;
-                } else {
-                    if (!editWithSnippet) {
-                        lineDelta = (indel.newText.match(/\n/g) || []).length - (indel.range.end.line - indel.range.start.line);
-                    }
-                    builder.replace(indel.range, indel.newText);
-                }
+    let selection: vscode.Selection | undefined = undefined;
+    let lineDelta = 0;
+    await editor.edit((builder) => {
+        for (const indel of edits) {
+            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;
+                selection = new vscode.Selection(
+                    new vscode.Position(startLine, startColumn),
+                    new vscode.Position(startLine, endColumn),
+                );
+                builder.replace(indel.range, newText);
+            } else {
+                lineDelta = countLines(indel.newText) - (indel.range.end.line - indel.range.start.line);
+                builder.replace(indel.range, indel.newText);
             }
-        });
-        if (editWithSnippet) {
-            const snip = editWithSnippet as vscode.TextEdit;
-            const range = snip.range.with(
-                snip.range.start.with(snip.range.start.line + lineDelta),
-                snip.range.end.with(snip.range.end.line + lineDelta),
-            );
-            await editor.insertSnippet(new vscode.SnippetString(snip.newText), range);
         }
-    };
+    });
+    if (selection) editor.selection = selection;
+}
+
+function parseSnippet(snip: string): [string, [number, number]] | undefined {
+    const m = snip.match(/\$(0|\{0:([^}]*)\})/);
+    if (!m) return undefined;
+    const placeholder = m[2] ?? "";
+    const range: [number, number] = [m.index!!, placeholder.length];
+    const insert = snip.replace(m[0], placeholder);
+    return [insert, range];
+}
+
+function countLines(text: string): number {
+    return (text.match(/\n/g) || []).length;
 }
diff --git a/editors/code/src/commands/on_enter.ts b/editors/code/src/commands/on_enter.ts
index 285849db7..a7871c31e 100644
--- a/editors/code/src/commands/on_enter.ts
+++ b/editors/code/src/commands/on_enter.ts
@@ -1,8 +1,8 @@
 import * as vscode from 'vscode';
 import * as ra from '../rust-analyzer-api';
 
-import { applySourceChange } from '../source_change';
 import { Cmd, Ctx } from '../ctx';
+import { applySnippetWorkspaceEdit } from '.';
 
 async function handleKeypress(ctx: Ctx) {
     const editor = ctx.activeRustEditor;
@@ -21,7 +21,8 @@ async function handleKeypress(ctx: Ctx) {
     });
     if (!change) return false;
 
-    await applySourceChange(ctx, change);
+    const workspaceEdit = client.protocol2CodeConverter.asWorkspaceEdit(change);
+    await applySnippetWorkspaceEdit(workspaceEdit);
     return true;
 }
 
diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index ac3bb365e..8b0a9d870 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -91,7 +91,7 @@ export async function activate(context: vscode.ExtensionContext) {
     ctx.registerCommand('debugSingle', commands.debugSingle);
     ctx.registerCommand('showReferences', commands.showReferences);
     ctx.registerCommand('applySourceChange', commands.applySourceChange);
-    ctx.registerCommand('applySnippetWorkspaceEdit', commands.applySnippetWorkspaceEdit);
+    ctx.registerCommand('applySnippetWorkspaceEdit', commands.applySnippetWorkspaceEditCommand);
     ctx.registerCommand('selectAndApplySourceChange', commands.selectAndApplySourceChange);
 
     ctx.pushCleanup(activateTaskProvider(workspaceFolder));
diff --git a/editors/code/src/rust-analyzer-api.ts b/editors/code/src/rust-analyzer-api.ts
index 400ac3714..3b83b10e3 100644
--- a/editors/code/src/rust-analyzer-api.ts
+++ b/editors/code/src/rust-analyzer-api.ts
@@ -69,7 +69,7 @@ export interface JoinLinesParams {
 export const joinLines = request<JoinLinesParams, SourceChange>("joinLines");
 
 
-export const onEnter = request<lc.TextDocumentPositionParams, Option<SourceChange>>("onEnter");
+export const onEnter = request<lc.TextDocumentPositionParams, Option<lc.WorkspaceEdit>>("onEnter");
 
 export interface RunnablesParams {
     textDocument: lc.TextDocumentIdentifier;