4 files changed, 44 insertions, 25 deletions

diff --git a/editors/code/src/extension.ts b/editors/code/src/extension.ts
index f1bc0b457..44e74f4cc 100644
--- a/editors/code/src/extension.ts
+++ b/editors/code/src/extension.ts
@@ -1,8 +1,10 @@
 import * as vscode from 'vscode';
+import * as lc from 'vscode-languageclient';
 
 import * as commands from './commands';
 import { TextDocumentContentProvider } from './commands/syntaxTree';
 import * as events from './events';
+import * as notifications from './notifications';
 import { Server } from './server';
 
 export function activate(context: vscode.ExtensionContext) {
@@ -14,6 +16,7 @@ export function activate(context: vscode.ExtensionContext) {
         disposeOnDeactivation(vscode.commands.registerCommand(name, f));
     }
 
+    // Commands are requests from vscode to the language server
     registerCommand('ra-lsp.syntaxTree', commands.syntaxTree.handle);
     registerCommand('ra-lsp.extendSelection', commands.extendSelection.handle);
     registerCommand('ra-lsp.matchingBrace', commands.matchingBrace.handle);
@@ -22,19 +25,27 @@ export function activate(context: vscode.ExtensionContext) {
     registerCommand('ra-lsp.run', commands.runnables.handle);
     registerCommand('ra-lsp.applySourceChange', commands.applySourceChange.handle);
 
+    // Notifications are events triggered by the language server
+    const allNotifications: Iterable<[string, lc.GenericNotificationHandler]> = [
+        ['m/publishDecorations', notifications.publishDecorations.handle],
+    ];
+
+    // The events below are plain old javascript events, triggered and handled by vscode
+    vscode.window.onDidChangeActiveTextEditor(events.changeActiveTextEditor.handle);
+
     const textDocumentContentProvider = new TextDocumentContentProvider();
     disposeOnDeactivation(vscode.workspace.registerTextDocumentContentProvider(
         'ra-lsp',
         textDocumentContentProvider,
     ));
 
-    Server.start();
-
     vscode.workspace.onDidChangeTextDocument(
         events.changeTextDocument.createHandler(textDocumentContentProvider),
         null,
         context.subscriptions);
-    vscode.window.onDidChangeActiveTextEditor(events.changeActiveTextEditor.handle);
+
+    // Start the language server, finally!
+    Server.start(allNotifications);
 }
 
 export function deactivate(): Thenable<void> {
diff --git a/editors/code/src/notifications/index.ts b/editors/code/src/notifications/index.ts
new file mode 100644
index 000000000..c56576865
--- /dev/null
+++ b/editors/code/src/notifications/index.ts
@@ -0,0 +1,5 @@
+import * as publishDecorations from './publish_decorations';
+
+export {
+    publishDecorations,
+};
diff --git a/editors/code/src/notifications/publish_decorations.ts b/editors/code/src/notifications/publish_decorations.ts
new file mode 100644
index 000000000..d8790386b
--- /dev/null
+++ b/editors/code/src/notifications/publish_decorations.ts
@@ -0,0 +1,20 @@
+import * as vscode from 'vscode';
+
+import { Decoration } from '../highlighting';
+import { Server } from '../server';
+
+export interface PublishDecorationsParams {
+    uri: string;
+    decorations: Decoration[];
+}
+
+export function handle(params: PublishDecorationsParams) {
+    const targetEditor = vscode.window.visibleTextEditors.find(
+        (editor) => editor.document.uri.toString() === params.uri,
+    );
+    if (!Server.config.highlightingOn || !targetEditor) { return; }
+    Server.highlighter.setHighlights(
+        targetEditor,
+        params.decorations,
+    );
+}
diff --git a/editors/code/src/server.ts b/editors/code/src/server.ts
index 325023e36..01fd80756 100644
--- a/editors/code/src/server.ts
+++ b/editors/code/src/server.ts
@@ -1,15 +1,14 @@
-import * as vscode from 'vscode';
 import * as lc from 'vscode-languageclient';
 
 import { Config } from './config';
-import { Decoration, Highlighter } from './highlighting';
+import { Highlighter } from './highlighting';
 
 export class Server {
     public static highlighter = new Highlighter();
     public static config = new Config();
     public static client: lc.LanguageClient;
 
-    public static start() {
+    public static start(notificationHandlers: Iterable<[string, lc.GenericNotificationHandler]>) {
         const run: lc.Executable = {
             command: 'ra_lsp_server',
             options: { cwd: '.' },
@@ -18,7 +17,6 @@ export class Server {
             run,
             debug: run,
         };
-
         const clientOptions: lc.LanguageClientOptions = {
             documentSelector: [{ scheme: 'file', language: 'rust' }],
         };
@@ -30,25 +28,10 @@ export class Server {
             clientOptions,
         );
         Server.client.onReady().then(() => {
-            Server.client.onNotification(
-                'm/publishDecorations',
-                (params: PublishDecorationsParams) => {
-                    const targetEditor = vscode.window.visibleTextEditors.find(
-                        (editor) => editor.document.uri.toString() === params.uri,
-                    );
-                    if (!Server.config.highlightingOn || !targetEditor) { return; }
-                    Server.highlighter.setHighlights(
-                        targetEditor,
-                        params.decorations,
-                    );
-                },
-            );
+            for (const [type, handler] of notificationHandlers) {
+                Server.client.onNotification(type, handler);
+            }
         });
         Server.client.start();
     }
 }
-
-interface PublishDecorationsParams {
-    uri: string;
-    decorations: Decoration[];
-}