5 files changed, 38 insertions, 45 deletions

diff --git a/editors/code/src/ctx.ts b/editors/code/src/ctx.ts
index 693ce05ed..0e62a3a85 100644
--- a/editors/code/src/ctx.ts
+++ b/editors/code/src/ctx.ts
@@ -84,6 +84,11 @@ export class Ctx {
         }
         throw 'unreachable';
     }
+
+    onNotification(method: string, handler: lc.GenericNotificationHandler) {
+        this.client.onReady()
+            .then(() => this.client.onNotification(method, handler))
+    }
 }
 
 export type Cmd = (...args: any[]) => any;
diff --git a/editors/code/src/highlighting.ts b/editors/code/src/highlighting.ts
index d383d87ef..d4e961b5b 100644
--- a/editors/code/src/highlighting.ts
+++ b/editors/code/src/highlighting.ts
@@ -10,28 +10,26 @@ import { Ctx } from './ctx';
 export function activateHighlighting(ctx: Ctx) {
     const highlighter = new Highlighter(ctx);
 
-    ctx.client.onReady().then(() => {
-        ctx.client.onNotification(
-            'rust-analyzer/publishDecorations',
-            (params: PublishDecorationsParams) => {
-                if (!ctx.config.highlightingOn) return;
-
-                const targetEditor = vscode.window.visibleTextEditors.find(
-                    editor => {
-                        const unescapedUri = unescape(
-                            editor.document.uri.toString(),
-                        );
-                        // Unescaped URI looks like:
-                        // file:///c:/Workspace/ra-test/src/main.rs
-                        return unescapedUri === params.uri;
-                    },
-                );
-                if (!targetEditor) return;
-
-                highlighter.setHighlights(targetEditor, params.decorations);
-            },
-        );
-    });
+    ctx.onNotification(
+        'rust-analyzer/publishDecorations',
+        (params: PublishDecorationsParams) => {
+            if (!ctx.config.highlightingOn) return;
+
+            const targetEditor = vscode.window.visibleTextEditors.find(
+                editor => {
+                    const unescapedUri = unescape(
+                        editor.document.uri.toString(),
+                    );
+                    // Unescaped URI looks like:
+                    // file:///c:/Workspace/ra-test/src/main.rs
+                    return unescapedUri === params.uri;
+                },
+            );
+            if (!targetEditor) return;
+
+            highlighter.setHighlights(targetEditor, params.decorations);
+        },
+    );
 
     vscode.workspace.onDidChangeConfiguration(
         _ => highlighter.removeHighlights(),
diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index 0c4abdac8..511f17ca4 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -1,9 +1,8 @@
 import * as vscode from 'vscode';
-import * as lc from 'vscode-languageclient';
 
 import * as commands from './commands';
 import { activateInlayHints } from './inlay_hints';
-import { StatusDisplay } from './status_display';
+import { activateStatusDisplay } from './status_display';
 import { Server } from './server';
 import { Ctx } from './ctx';
 import { activateHighlighting } from './highlighting';
@@ -32,18 +31,7 @@ export async function activate(context: vscode.ExtensionContext) {
         ctx.overrideCommand('type', commands.onEnter);
     }
 
-    const watchStatus = new StatusDisplay(ctx.config.cargoWatchOptions.command);
-    ctx.pushCleanup(watchStatus);
-
-    // Notifications are events triggered by the language server
-    const allNotifications: [string, lc.GenericNotificationHandler][] = [
-        [
-            '$/progress',
-            params => watchStatus.handleProgressNotification(params),
-        ],
-    ];
-
-    const startServer = () => Server.start(allNotifications);
+    const startServer = () => Server.start();
     const reloadCommand = () => reloadServer(startServer);
 
     vscode.commands.registerCommand('rust-analyzer.reload', reloadCommand);
@@ -55,6 +43,7 @@ export async function activate(context: vscode.ExtensionContext) {
         vscode.window.showErrorMessage(e.message);
     }
 
+    activateStatusDisplay(ctx);
     activateHighlighting(ctx);
 
     if (ctx.config.displayInlayHints) {
diff --git a/editors/code/src/server.ts b/editors/code/src/server.ts
index 2bb21da6b..5dc8a36bd 100644
--- a/editors/code/src/server.ts
+++ b/editors/code/src/server.ts
@@ -15,9 +15,7 @@ export class Server {
     public static config = new Config();
     public static client: lc.LanguageClient;
 
-    public static async start(
-        notificationHandlers: Iterable<[string, lc.GenericNotificationHandler]>,
-    ) {
+    public static async start() {
         // '.' Is the fallback if no folder is open
         // TODO?: Workspace folders support Uri's (eg: file://test.txt). It might be a good idea to test if the uri points to a file.
         let folder: string = '.';
@@ -92,11 +90,6 @@ export class Server {
             },
         };
         Server.client.registerProposedFeatures();
-        Server.client.onReady().then(() => {
-            for (const [type, handler] of notificationHandlers) {
-                Server.client.onNotification(type, handler);
-            }
-        });
         Server.client.start();
     }
 }
diff --git a/editors/code/src/status_display.ts b/editors/code/src/status_display.ts
index 48cf0655b..e3719075b 100644
--- a/editors/code/src/status_display.ts
+++ b/editors/code/src/status_display.ts
@@ -1,8 +1,16 @@
 import * as vscode from 'vscode';
 
+import { Ctx } from './ctx';
+
 const spinnerFrames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
 
-export class StatusDisplay implements vscode.Disposable {
+export function activateStatusDisplay(ctx: Ctx) {
+    const statusDisplay = new StatusDisplay(ctx.config.cargoWatchOptions.command);
+    ctx.pushCleanup(statusDisplay);
+    ctx.onNotification('$/progress', params => statusDisplay.handleProgressNotification(params));
+}
+
+class StatusDisplay implements vscode.Disposable {
     packageName?: string;
 
     private i = 0;