feat: Extend vscode 'run' command with optional mode argument for running test(s) or bin at keyboard cursor

3 files changed, 70 insertions, 4 deletions

diff --git a/src/tools/rust-analyzer/editors/code/src/commands.ts b/src/tools/rust-analyzer/editors/code/src/commands.ts
index 3ac1a933d9e..25b30013fa1 100644
--- a/src/tools/rust-analyzer/editors/code/src/commands.ts
+++ b/src/tools/rust-analyzer/editors/code/src/commands.ts
@@ -1114,11 +1114,11 @@ export function applySnippetWorkspaceEditCommand(_ctx: CtxInit): Cmd {
     };
 }
 
-export function run(ctx: CtxInit): Cmd {
+export function run(ctx: CtxInit, mode?: "cursor"): Cmd {
     let prevRunnable: RunnableQuickPick | undefined;
 
     return async () => {
-        const item = await selectRunnable(ctx, prevRunnable);
+        const item = await selectRunnable(ctx, prevRunnable, false, true, mode);
         if (!item) return;
 
         item.detail = "rerun";
diff --git a/src/tools/rust-analyzer/editors/code/src/ctx.ts b/src/tools/rust-analyzer/editors/code/src/ctx.ts
index e55754fb9f0..0fb0a1f17cb 100644
--- a/src/tools/rust-analyzer/editors/code/src/ctx.ts
+++ b/src/tools/rust-analyzer/editors/code/src/ctx.ts
@@ -1,6 +1,7 @@
 import * as vscode from "vscode";
 import type * as lc from "vscode-languageclient/node";
 import * as ra from "./lsp_ext";
+import * as commands from "./commands";
 
 import { Config, prepareVSCodeConfig } from "./config";
 import { createClient } from "./client";
@@ -462,9 +463,17 @@ export class Ctx implements RustAnalyzerExtensionApi {
         for (const [name, factory] of Object.entries(this.commandFactories)) {
             const fullName = `rust-analyzer.${name}`;
             let callback;
+
             if (isClientRunning(this)) {
-                // we asserted that `client` is defined
-                callback = factory.enabled(this);
+                if (name === "run") {
+                    // Special case: support optional argument for `run`
+                    callback = (mode?: "cursor") => {
+                        return commands.run(this, mode)();
+                    };
+                } else {
+                    // we asserted that `client` is defined
+                    callback = factory.enabled(this);
+                }
             } else if (factory.disabled) {
                 callback = factory.disabled(this);
             } else {
diff --git a/src/tools/rust-analyzer/editors/code/src/run.ts b/src/tools/rust-analyzer/editors/code/src/run.ts
index 40027cc7c85..95166c427b2 100644
--- a/src/tools/rust-analyzer/editors/code/src/run.ts
+++ b/src/tools/rust-analyzer/editors/code/src/run.ts
@@ -18,10 +18,15 @@ export async function selectRunnable(
     prevRunnable?: RunnableQuickPick,
     debuggeeOnly = false,
     showButtons: boolean = true,
+    mode?: "cursor",
 ): Promise<RunnableQuickPick | undefined> {
     const editor = ctx.activeRustEditor ?? ctx.activeCargoTomlEditor;
     if (!editor) return;
 
+    if (mode === "cursor") {
+        return selectRunnableAtCursor(ctx, editor, prevRunnable);
+    }
+
     // show a placeholder while we get the runnables from the server
     const quickPick = vscode.window.createQuickPick();
     quickPick.title = "Select Runnable";
@@ -54,6 +59,58 @@ export async function selectRunnable(
     );
 }
 
+async function selectRunnableAtCursor(
+    ctx: CtxInit,
+    editor: RustEditor,
+    prevRunnable?: RunnableQuickPick,
+): Promise<RunnableQuickPick | undefined> {
+    const runnableQuickPicks = await getRunnables(ctx.client, editor, prevRunnable, false);
+    let runnableQuickPickAtCursor = null;
+    const cursorPosition = ctx.client.code2ProtocolConverter.asPosition(editor.selection.active);
+    for (const runnableQuickPick of runnableQuickPicks) {
+        if (!runnableQuickPick.runnable.location?.targetRange) {
+            continue;
+        }
+        const runnableQuickPickRange = runnableQuickPick.runnable.location.targetRange;
+        if (
+            runnableQuickPickAtCursor?.runnable?.location?.targetRange != null &&
+            rangeContainsOtherRange(
+                runnableQuickPickRange,
+                runnableQuickPickAtCursor.runnable.location.targetRange,
+            )
+        ) {
+            continue;
+        }
+        if (rangeContainsPosition(runnableQuickPickRange, cursorPosition)) {
+            runnableQuickPickAtCursor = runnableQuickPick;
+        }
+    }
+    if (runnableQuickPickAtCursor == null) {
+        return;
+    }
+    return Promise.resolve(runnableQuickPickAtCursor);
+}
+
+function rangeContainsPosition(range: lc.Range, position: lc.Position): boolean {
+    return (
+        (position.line > range.start.line ||
+            (position.line === range.start.line && position.character >= range.start.character)) &&
+        (position.line < range.end.line ||
+            (position.line === range.end.line && position.character <= range.end.character))
+    );
+}
+
+function rangeContainsOtherRange(range: lc.Range, otherRange: lc.Range) {
+    return (
+        (range.start.line < otherRange.start.line ||
+            (range.start.line === otherRange.start.line &&
+                range.start.character <= otherRange.start.character)) &&
+        (range.end.line > otherRange.end.line ||
+            (range.end.line === otherRange.end.line &&
+                range.end.character >= otherRange.end.character))
+    );
+}
+
 export class RunnableQuickPick implements vscode.QuickPickItem {
     public label: string;
     public description?: string | undefined;