Simplify TS reload logic

Fixes #3164

6 files changed, 51 insertions, 23 deletions

diff --git a/editors/code/src/commands/index.ts b/editors/code/src/commands/index.ts
index b5ebec117f3..d05f40d67af 100644
--- a/editors/code/src/commands/index.ts
+++ b/editors/code/src/commands/index.ts
@@ -51,10 +51,3 @@ export function selectAndApplySourceChange(ctx: Ctx): Cmd {
         }
     };
 }
-
-export function reload(ctx: Ctx): Cmd {
-    return async () => {
-        vscode.window.showInformationMessage('Reloading rust-analyzer...');
-        await ctx.restartServer();
-    };
-}
diff --git a/editors/code/src/ctx.ts b/editors/code/src/ctx.ts
index ff6245f7899..1eff88df2a4 100644
--- a/editors/code/src/ctx.ts
+++ b/editors/code/src/ctx.ts
@@ -1,5 +1,6 @@
 import * as vscode from 'vscode';
 import * as lc from 'vscode-languageclient';
+import { strict as assert } from "assert";
 
 import { Config } from './config';
 import { createClient } from './client';
@@ -16,19 +17,16 @@ export class Ctx {
     // on the event loop to get a better picture of what we can do here)
     client: lc.LanguageClient | null = null;
     private extCtx: vscode.ExtensionContext;
-    private onDidRestartHooks: Array<(client: lc.LanguageClient) => void> = [];
+    private onStartHooks: Array<(client: lc.LanguageClient) => void> = [];
 
     constructor(extCtx: vscode.ExtensionContext) {
         this.config = new Config(extCtx);
         this.extCtx = extCtx;
     }
 
-    async restartServer() {
-        const old = this.client;
-        if (old) {
-            await old.stop();
-        }
-        this.client = null;
+    async startServer() {
+        assert(this.client == null);
+
         const client = await createClient(this.config);
         if (!client) {
             throw new Error(
@@ -41,7 +39,7 @@ export class Ctx {
         await client.onReady();
 
         this.client = client;
-        for (const hook of this.onDidRestartHooks) {
+        for (const hook of this.onStartHooks) {
             hook(client);
         }
     }
@@ -72,8 +70,13 @@ export class Ctx {
         this.extCtx.subscriptions.push(d);
     }
 
-    onDidRestart(hook: (client: lc.LanguageClient) => void) {
-        this.onDidRestartHooks.push(hook);
+    onStart(hook: (client: lc.LanguageClient) => void) {
+        const client = this.client;
+        if (client == null) {
+            this.onStartHooks.push(hook);
+        } else {
+            hook(client)
+        }
     }
 }
 
diff --git a/editors/code/src/highlighting.ts b/editors/code/src/highlighting.ts
index 4fbbe3ddc5c..f693fb8ba78 100644
--- a/editors/code/src/highlighting.ts
+++ b/editors/code/src/highlighting.ts
@@ -7,7 +7,7 @@ import { Ctx, sendRequestWithRetry } from './ctx';
 
 export function activateHighlighting(ctx: Ctx) {
     const highlighter = new Highlighter(ctx);
-    ctx.onDidRestart(client => {
+    ctx.onStart(client => {
         client.onNotification(
             'rust-analyzer/publishDecorations',
             (params: PublishDecorationsParams) => {
diff --git a/editors/code/src/inlay_hints.ts b/editors/code/src/inlay_hints.ts
index 3896878cda5..9e400fabec0 100644
--- a/editors/code/src/inlay_hints.ts
+++ b/editors/code/src/inlay_hints.ts
@@ -27,9 +27,15 @@ export function activateInlayHints(ctx: Ctx) {
         ctx.subscriptions
     );
 
+    ctx.pushCleanup({
+        dispose() {
+            hintsUpdater.clear()
+        }
+    })
+
     // We pass async function though it will not be awaited when called,
     // thus Promise rejections won't be handled, but this should never throw in fact...
-    ctx.onDidRestart(async _ => hintsUpdater.setEnabled(ctx.config.displayInlayHints));
+    ctx.onStart(async _ => hintsUpdater.setEnabled(ctx.config.displayInlayHints));
 }
 
 interface InlayHintsParams {
@@ -61,16 +67,23 @@ class HintsUpdater {
 
     constructor(ctx: Ctx) {
         this.ctx = ctx;
-        this.enabled = ctx.config.displayInlayHints;
+        this.enabled = false;
     }
 
     async setEnabled(enabled: boolean): Promise<void> {
+        console.log({ enabled, prev: this.enabled });
+
         if (this.enabled == enabled) return;
         this.enabled = enabled;
 
         if (this.enabled) {
             return await this.refresh();
+        } else {
+            return this.clear();
         }
+    }
+
+    clear() {
         this.allEditors.forEach(it => {
             this.setTypeDecorations(it, []);
             this.setParameterDecorations(it, []);
@@ -79,6 +92,8 @@ class HintsUpdater {
 
     async refresh() {
         if (!this.enabled) return;
+        console.log("freshin!");
+
         await Promise.all(this.allEditors.map(it => this.refreshEditor(it)));
     }
 
diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index 5a99e96f0e5..ec488c34055 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -11,6 +11,23 @@ let ctx: Ctx | undefined;
 export async function activate(context: vscode.ExtensionContext) {
     ctx = new Ctx(context);
 
+    ctx.registerCommand('reload', (ctx) => {
+        return async () => {
+            vscode.window.showInformationMessage('Reloading rust-analyzer...');
+            // @DanTup maneuver
+            // https://github.com/microsoft/vscode/issues/45774#issuecomment-373423895
+            await deactivate()
+            for (const sub of ctx.subscriptions) {
+                try {
+                    sub.dispose();
+                } catch (e) {
+                    console.error(e);
+                }
+            }
+            await activate(context)
+        }
+    })
+
     // Commands which invokes manually via command palette, shortcut, etc.
     ctx.registerCommand('analyzerStatus', commands.analyzerStatus);
     ctx.registerCommand('collectGarbage', commands.collectGarbage);
@@ -20,7 +37,6 @@ export async function activate(context: vscode.ExtensionContext) {
     ctx.registerCommand('syntaxTree', commands.syntaxTree);
     ctx.registerCommand('expandMacro', commands.expandMacro);
     ctx.registerCommand('run', commands.run);
-    ctx.registerCommand('reload', commands.reload);
     ctx.registerCommand('onEnter', commands.onEnter);
     ctx.registerCommand('ssr', commands.ssr)
 
@@ -38,7 +54,7 @@ export async function activate(context: vscode.ExtensionContext) {
     //
     // This a horribly, horribly wrong way to deal with this problem.
     try {
-        await ctx.restartServer();
+        await ctx.startServer();
     } catch (e) {
         vscode.window.showErrorMessage(e.message);
     }
@@ -47,4 +63,5 @@ export async function activate(context: vscode.ExtensionContext) {
 
 export async function deactivate() {
     await ctx?.client?.stop();
+    ctx = undefined;
 }
diff --git a/editors/code/src/status_display.ts b/editors/code/src/status_display.ts
index 993e79d7036..326b5217b5b 100644
--- a/editors/code/src/status_display.ts
+++ b/editors/code/src/status_display.ts
@@ -9,7 +9,7 @@ const spinnerFrames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '
 export function activateStatusDisplay(ctx: Ctx) {
     const statusDisplay = new StatusDisplay(ctx.config.cargoWatchOptions.command);
     ctx.pushCleanup(statusDisplay);
-    ctx.onDidRestart(client => ctx.pushCleanup(client.onProgress(
+    ctx.onStart(client => ctx.pushCleanup(client.onProgress(
         WorkDoneProgress.type,
         'rustAnalyzer/cargoWatcher',
         params => statusDisplay.handleProgressNotification(params)