[clangd] Cleans up the semantic highlighting resources if clangd stops.

Summary: Disposes of the vscode listeners when clangd crashes and reuses the old highlighter when it restarts. The reason for reusing the highlighter is because this way the highlightings will not disappear as we won't have to dispose of them.

Reviewers: hokein, ilya-biryukov

Subscribers: MaskRay, jkorous, arphaman, kadircet, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D66743

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@370202 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/clangd/clients/clangd-vscode/src/extension.ts b/clangd/clients/clangd-vscode/src/extension.ts
index dc187ee..f06b2f1 100644
--- a/clangd/clients/clangd-vscode/src/extension.ts
+++ b/clangd/clients/clangd-vscode/src/extension.ts
@@ -112,14 +112,6 @@
   const semanticHighlightingFeature =
       new semanticHighlighting.SemanticHighlightingFeature();
   clangdClient.registerFeature(semanticHighlightingFeature);
-  // The notification handler must be registered after the client is ready or
-  // the client will crash.
-  clangdClient.onReady().then(
-      () => clangdClient.onNotification(
-          semanticHighlighting.NotificationType,
-          semanticHighlightingFeature.handleNotification.bind(
-              semanticHighlightingFeature)));
-
   console.log('Clang Language Server is now active!');
   context.subscriptions.push(clangdClient.start());
   context.subscriptions.push(vscode.commands.registerCommand(
@@ -149,9 +141,14 @@
       clangdClient.onNotification(
           'textDocument/clangd.fileStatus',
           (fileStatus) => { status.onFileUpdated(fileStatus); });
+      clangdClient.onNotification(
+          semanticHighlighting.NotificationType,
+          semanticHighlightingFeature.handleNotification.bind(
+              semanticHighlightingFeature));
     } else if (newState == vscodelc.State.Stopped) {
       // Clear all cached statuses when clangd crashes.
       status.clear();
+      semanticHighlightingFeature.dispose();
     }
   })
   // An empty place holder for the activate command, otherwise we'll get an
diff --git a/clangd/clients/clangd-vscode/src/semantic-highlighting.ts b/clangd/clients/clangd-vscode/src/semantic-highlighting.ts
index ab9f45f..2ac700e 100644
--- a/clangd/clients/clangd-vscode/src/semantic-highlighting.ts
+++ b/clangd/clients/clangd-vscode/src/semantic-highlighting.ts
@@ -56,6 +56,8 @@
   scopeLookupTable: string[][];
   // The object that applies the highlightings clangd sends.
   highlighter: Highlighter;
+  // Any disposables that should be cleaned up when clangd crashes.
+  private subscriptions: vscode.Disposable[] = [];
   fillClientCapabilities(capabilities: vscodelc.ClientCapabilities) {
     // Extend the ClientCapabilities type and add semantic highlighting
     // capability to the object.
@@ -88,14 +90,15 @@
     // otherwise it could try to update the themeRuleMatcher without the
     // highlighter being created.
     this.highlighter = new Highlighter(this.scopeLookupTable);
+    this.subscriptions.push(vscode.Disposable.from(this.highlighter));
     this.loadCurrentTheme();
     // Event handling for handling with TextDocuments/Editors lifetimes.
-    vscode.window.onDidChangeVisibleTextEditors(
+    this.subscriptions.push(vscode.window.onDidChangeVisibleTextEditors(
         (editors: vscode.TextEditor[]) =>
             editors.forEach((e) => this.highlighter.applyHighlights(
-                                e.document.uri.toString())));
-    vscode.workspace.onDidCloseTextDocument(
-        (doc) => this.highlighter.removeFileHighlightings(doc.uri.toString()));
+                                e.document.uri.toString()))));
+    this.subscriptions.push(vscode.workspace.onDidCloseTextDocument(
+        (doc) => this.highlighter.removeFileHighlightings(doc.uri.toString())));
   }
 
   handleNotification(params: SemanticHighlightingParams) {
@@ -103,6 +106,11 @@
         (line) => ({line : line.line, tokens : decodeTokens(line.tokens)}));
     this.highlighter.highlight(params.textDocument.uri, lines);
   }
+  // Disposes of all disposable resources used by this object.
+  public dispose() {
+    this.subscriptions.forEach((d) => d.dispose());
+    this.subscriptions = [];
+  }
 }
 
 // Converts a string of base64 encoded tokens into the corresponding array of
@@ -138,6 +146,13 @@
   constructor(scopeLookupTable: string[][]) {
     this.scopeLookupTable = scopeLookupTable;
   }
+  public dispose() {
+    this.files.clear();
+    this.decorationTypes.forEach((t) => t.dispose());
+    // Dispose must not be not called multiple times if initialize is
+    // called again.
+    this.decorationTypes = [];
+  }
   // This function must be called at least once or no highlightings will be
   // done. Sets the theme that is used when highlighting. Also triggers a
   // recolorization for all current highlighters. Should be called whenever the
@@ -174,6 +189,27 @@
     this.applyHighlights(fileUri);
   }
 
+  // Applies all the highlightings currently stored for a file with fileUri.
+  public applyHighlights(fileUri: string) {
+    if (!this.files.has(fileUri))
+      // There are no highlightings for this file, must return early or will get
+      // out of bounds when applying the decorations below.
+      return;
+    if (!this.decorationTypes.length)
+      // Can't apply any decorations when there is no theme loaded.
+      return;
+    // This must always do a full re-highlighting due to the fact that
+    // TextEditorDecorationType are very expensive to create (which makes
+    // incremental updates infeasible). For this reason one
+    // TextEditorDecorationType is used per scope.
+    const ranges = this.getDecorationRanges(fileUri);
+    vscode.window.visibleTextEditors.forEach((e) => {
+      if (e.document.uri.toString() !== fileUri)
+        return;
+      this.decorationTypes.forEach((d, i) => e.setDecorations(d, ranges[i]));
+    });
+  }
+
   // Called when a text document is closed. Removes any highlighting entries for
   // the text document that was closed.
   public removeFileHighlightings(fileUri: string) {
@@ -207,27 +243,6 @@
     });
     return decorations;
   }
-
-  // Applies all the highlightings currently stored for a file with fileUri.
-  public applyHighlights(fileUri: string) {
-    if (!this.files.has(fileUri))
-      // There are no highlightings for this file, must return early or will get
-      // out of bounds when applying the decorations below.
-      return;
-    if (!this.decorationTypes.length)
-      // Can't apply any decorations when there is no theme loaded.
-      return;
-    // This must always do a full re-highlighting due to the fact that
-    // TextEditorDecorationType are very expensive to create (which makes
-    // incremental updates infeasible). For this reason one
-    // TextEditorDecorationType is used per scope.
-    const ranges = this.getDecorationRanges(fileUri);
-    vscode.window.visibleTextEditors.forEach((e) => {
-      if (e.document.uri.toString() !== fileUri)
-        return;
-      this.decorationTypes.forEach((d, i) => e.setDecorations(d, ranges[i]));
-    });
-  }
 }
 
 // A rule for how to color TextMate scopes.
diff --git a/clangd/clients/clangd-vscode/test/semantic-highlighting.test.ts b/clangd/clients/clangd-vscode/test/semantic-highlighting.test.ts
index 0f099bd..6bc3c45 100644
--- a/clangd/clients/clangd-vscode/test/semantic-highlighting.test.ts
+++ b/clangd/clients/clangd-vscode/test/semantic-highlighting.test.ts
@@ -107,7 +107,8 @@
     highlighter.highlight('file1', []);
     assert.deepEqual(highlighter.applicationUriHistory, [ 'file1' ]);
     highlighter.initialize(tm);
-    assert.deepEqual(highlighter.applicationUriHistory, [ 'file1', 'file1', 'file2' ]);
+    assert.deepEqual(highlighter.applicationUriHistory,
+                     [ 'file1', 'file1', 'file2' ]);
     // Groups decorations into the scopes used.
     let highlightingsInLine: semanticHighlighting.SemanticHighlightingLine[] = [
       {