Changes for debug restarting

4 files changed, 63 insertions, 1 deletions

diff --git a/src/tools/rust-analyzer/editors/code/package.json b/src/tools/rust-analyzer/editors/code/package.json
index bf9c4a366d4..1406338a6cd 100644
--- a/src/tools/rust-analyzer/editors/code/package.json
+++ b/src/tools/rust-analyzer/editors/code/package.json
@@ -520,6 +520,11 @@
                         "type": "boolean",
                         "default": false
                     },
+                    "rust-analyzer.debug.buildBeforeRestart": {
+                        "markdownDescription": "Whether to rebuild the project modules before debugging the same test again",
+                        "type": "boolean",
+                        "default": false
+                    },
                     "rust-analyzer.debug.engineSettings": {
                         "type": "object",
                         "default": {},
diff --git a/src/tools/rust-analyzer/editors/code/src/config.ts b/src/tools/rust-analyzer/editors/code/src/config.ts
index dc0165df71e..6c5c1173da6 100644
--- a/src/tools/rust-analyzer/editors/code/src/config.ts
+++ b/src/tools/rust-analyzer/editors/code/src/config.ts
@@ -303,6 +303,7 @@ export class Config {
             engine: this.get<string>("debug.engine"),
             engineSettings: this.get<object>("debug.engineSettings") ?? {},
             openDebugPane: this.get<boolean>("debug.openDebugPane"),
+            buildBeforeRestart: this.get<boolean>("debug.buildBeforeRestart"),
             sourceFileMap: sourceFileMap,
         };
     }
diff --git a/src/tools/rust-analyzer/editors/code/src/debug.ts b/src/tools/rust-analyzer/editors/code/src/debug.ts
index 3aae0f9ce6e..b3f1b056a8b 100644
--- a/src/tools/rust-analyzer/editors/code/src/debug.ts
+++ b/src/tools/rust-analyzer/editors/code/src/debug.ts
@@ -5,12 +5,15 @@ import type * as ra from "./lsp_ext";
 
 import { Cargo } from "./toolchain";
 import type { Ctx } from "./ctx";
-import { prepareEnv } from "./run";
+import { createTaskFromRunnable, prepareEnv } from "./run";
 import { execute, isCargoRunnableArgs, unwrapUndefinable } from "./util";
 import type { Config } from "./config";
 
 const debugOutput = vscode.window.createOutputChannel("Debug");
 
+// Here we want to keep track on everything that's currently running
+const activeDebugSessionIds: string[] = [];
+
 export async function makeDebugConfig(ctx: Ctx, runnable: ra.Runnable): Promise<void> {
     const scope = ctx.activeRustEditor?.document.uri;
     if (!scope) return;
@@ -45,6 +48,8 @@ export async function startDebugSession(ctx: Ctx, runnable: ra.Runnable): Promis
     const wsLaunchSection = vscode.workspace.getConfiguration("launch");
     const configurations = wsLaunchSection.get<any[]>("configurations") || [];
 
+    // The runnable label is the name of the test with the "test prefix"
+    // e.g. test test_feature_x
     const index = configurations.findIndex((c) => c.name === runnable.label);
     if (-1 !== index) {
         debugConfig = configurations[index];
@@ -359,3 +364,49 @@ function quote(xs: string[]) {
         })
         .join(" ");
 }
+
+async function recompileTestFromDebuggingSession(session: vscode.DebugSession, ctx: Ctx) {
+    const { cwd, args: sessionArgs }: vscode.DebugConfiguration = session.configuration;
+
+    const args: ra.CargoRunnableArgs = {
+        cwd: cwd,
+        cargoArgs: ["test", "--no-run", "--test", "lib"],
+
+        // The first element of the debug configuration args is the test path e.g. "test_bar::foo::test_a::test_b"
+        executableArgs: sessionArgs,
+    };
+    const runnable: ra.Runnable = {
+        kind: "cargo",
+        label: "compile-test",
+        args,
+    };
+    const task: vscode.Task = await createTaskFromRunnable(runnable, ctx.config);
+
+    // It is not needed to call the language server, since the test path is already resolved in the
+    // configuration option. We can simply call a debug configuration with the --no-run option to compile
+    await vscode.tasks.executeTask(task);
+}
+
+export function initializeDebugSessionTrackingAndRebuild(ctx: Ctx) {
+    vscode.debug.onDidStartDebugSession((session: vscode.DebugSession) => {
+        if (!activeDebugSessionIds.includes(session.id)) {
+            activeDebugSessionIds.push(session.id);
+        }
+    });
+
+    vscode.debug.onDidTerminateDebugSession(async (session: vscode.DebugSession) => {
+        // The id of the session will be the same when pressing restart the restart button
+        if (activeDebugSessionIds.find((s) => s === session.id)) {
+            await recompileTestFromDebuggingSession(session, ctx);
+        }
+        removeActiveSession(session);
+    });
+}
+
+function removeActiveSession(session: vscode.DebugSession) {
+    const activeSessionId = activeDebugSessionIds.findIndex((id) => id === session.id);
+
+    if (activeSessionId !== -1) {
+        activeDebugSessionIds.splice(activeSessionId, 1);
+    }
+}
diff --git a/src/tools/rust-analyzer/editors/code/src/main.ts b/src/tools/rust-analyzer/editors/code/src/main.ts
index 4769fdd864a..0ddc5619e99 100644
--- a/src/tools/rust-analyzer/editors/code/src/main.ts
+++ b/src/tools/rust-analyzer/editors/code/src/main.ts
@@ -6,6 +6,7 @@ import { type CommandFactory, Ctx, fetchWorkspace } from "./ctx";
 import * as diagnostics from "./diagnostics";
 import { activateTaskProvider } from "./tasks";
 import { setContextValue } from "./util";
+import { initializeDebugSessionTrackingAndRebuild } from "./debug";
 
 const RUST_PROJECT_CONTEXT_NAME = "inRustProject";
 
@@ -102,6 +103,10 @@ async function activateServer(ctx: Ctx): Promise<RustAnalyzerExtensionApi> {
         ctx.subscriptions,
     );
 
+    if (ctx.config.debug.buildBeforeRestart) {
+        initializeDebugSessionTrackingAndRebuild(ctx);
+    }
+
     await ctx.start();
     return ctx;
 }