From efb2f9f2add20ea73e23d41e7c9464cc7aa50fe9 Mon Sep 17 00:00:00 2001
From: Paul LeMarquand <plemarquand@apple.com>
Date: Mon, 15 Jul 2024 10:31:26 -0400
Subject: [PATCH] Add Re-Index Project Command (#964)

* Add Re-Index Project Command

Adds a command to trigger a re-indexing of the open project. Should only
be needed when background indexing is enabled and it gets out of sync.
While that is a bug, this can act as a temporary workaround.
---
 package.json                         |  9 +++++++++
 src/TestExplorer/LSPTestDiscovery.ts |  2 --
 src/WorkspaceContext.ts              |  8 ++++++++
 src/commands.ts                      | 21 +++++++++++++++++++++
 src/contextKeys.ts                   |  7 +++++++
 src/sourcekit-lsp/lspExtensions.ts   |  4 ++++
 6 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/package.json b/package.json
index 11d24f644..bd6d70989 100644
--- a/package.json
+++ b/package.json
@@ -152,6 +152,11 @@
         "title": "Restart LSP Server",
         "category": "Swift"
       },
+      {
+        "command": "swift.reindexProject",
+        "title": "Re-Index Project",
+        "category": "Swift"
+      },
       {
         "command": "swift.switchPlatform",
         "title": "Select Target Platform...",
@@ -677,6 +682,10 @@
         {
           "command": "swift.attachDebugger",
           "when": "swift.lldbVSCodeAvailable"
+        },
+        {
+          "command": "swift.reindexProject",
+          "when": "swift.supportsReindexing"
         }
       ],
       "editor/context": [
diff --git a/src/TestExplorer/LSPTestDiscovery.ts b/src/TestExplorer/LSPTestDiscovery.ts
index 2b0d4fe9f..439c3f807 100644
--- a/src/TestExplorer/LSPTestDiscovery.ts
+++ b/src/TestExplorer/LSPTestDiscovery.ts
@@ -49,8 +49,6 @@ interface ILanguageClientManager {
  * these results.
  */
 export class LSPTestDiscovery {
-    private capCache = new Map<string, boolean>();
-
     constructor(private languageClient: ILanguageClientManager) {}
 
     /**
diff --git a/src/WorkspaceContext.ts b/src/WorkspaceContext.ts
index 92991353e..598992632 100644
--- a/src/WorkspaceContext.ts
+++ b/src/WorkspaceContext.ts
@@ -229,6 +229,14 @@ export class WorkspaceContext implements vscode.Disposable {
         } else {
             contextKeys.currentTargetType = undefined;
         }
+
+        // LSP can be configured per workspace to support reindexing
+        this.languageClientManager.useLanguageClient(async client => {
+            const experimentalCaps = client.initializeResult?.capabilities.experimental;
+            contextKeys.supportsReindexing =
+                experimentalCaps && experimentalCaps["workspace/triggerReindex"] !== undefined;
+        });
+
         setSnippetContextKey(this);
     }
 
diff --git a/src/commands.ts b/src/commands.ts
index 374b9854d..2f0d998bb 100644
--- a/src/commands.ts
+++ b/src/commands.ts
@@ -32,6 +32,7 @@ import { SwiftExecOperation, TaskOperation } from "./tasks/TaskQueue";
 import { SwiftProjectTemplate } from "./toolchain/toolchain";
 import { showToolchainSelectionQuickPick, showToolchainError } from "./ui/ToolchainSelection";
 import { captureDiagnostics } from "./commands/captureDiagnostics";
+import { reindexProjectRequest } from "./sourcekit-lsp/lspExtensions";
 
 /**
  * References:
@@ -656,6 +657,25 @@ function restartLSPServer(workspaceContext: WorkspaceContext): Promise<void> {
     return workspaceContext.languageClientManager.restart();
 }
 
+/** Request that the SourceKit-LSP server reindexes the workspace */
+function reindexProject(workspaceContext: WorkspaceContext): Promise<unknown> {
+    return workspaceContext.languageClientManager.useLanguageClient(async (client, token) => {
+        try {
+            return await client.sendRequest(reindexProjectRequest, {}, token);
+        } catch (err) {
+            const error = err as { code: number; message: string };
+            // methodNotFound, version of sourcekit-lsp is likely too old.
+            if (error.code === -32601) {
+                vscode.window.showWarningMessage(
+                    "The installed version of SourceKit-LSP does not support background indexing."
+                );
+            } else {
+                vscode.window.showWarningMessage(error.message);
+            }
+        }
+    });
+}
+
 /** Execute task and show UI while running */
 async function executeTaskWithUI(
     task: vscode.Task,
@@ -817,6 +837,7 @@ export function register(ctx: WorkspaceContext): vscode.Disposable[] {
         vscode.commands.registerCommand("swift.debugSnippet", () => debugSnippet(ctx)),
         vscode.commands.registerCommand("swift.runPluginTask", () => runPluginTask()),
         vscode.commands.registerCommand("swift.restartLSPServer", () => restartLSPServer(ctx)),
+        vscode.commands.registerCommand("swift.reindexProject", () => reindexProject(ctx)),
         vscode.commands.registerCommand("swift.insertFunctionComment", () =>
             insertFunctionComment(ctx)
         ),
diff --git a/src/contextKeys.ts b/src/contextKeys.ts
index a1d40e821..2cd037e08 100644
--- a/src/contextKeys.ts
+++ b/src/contextKeys.ts
@@ -80,6 +80,13 @@ const contextKeys = {
     set createNewProjectAvailable(value: boolean) {
         vscode.commands.executeCommand("setContext", "swift.createNewProjectAvailable", value);
     },
+
+    /**
+     * Whether the SourceKit-LSP server supports reindexing the workspace.
+     */
+    set supportsReindexing(value: boolean) {
+        vscode.commands.executeCommand("setContext", "swift.supportsReindexing", value);
+    },
 };
 
 export default contextKeys;
diff --git a/src/sourcekit-lsp/lspExtensions.ts b/src/sourcekit-lsp/lspExtensions.ts
index 6837a22f2..08a17137f 100644
--- a/src/sourcekit-lsp/lspExtensions.ts
+++ b/src/sourcekit-lsp/lspExtensions.ts
@@ -174,3 +174,7 @@ export const textDocumentTestsRequest = new langclient.RequestType<
     LSPTestItem[],
     unknown
 >("textDocument/tests");
+
+export const reindexProjectRequest = new langclient.RequestType<null, unknown, unknown>(
+    "workspace/triggerReindex"
+);