From e4a5720ea9817230a80a78c22d31ef5a97ec90c5 Mon Sep 17 00:00:00 2001
From: imkiva <imkiva@islovely.icu>
Date: Tue, 16 Nov 2021 22:44:18 +0800
Subject: [PATCH] HIGHLIGHT: workaroundly fix #15

---
 src/features.ts  |  7 +++++++
 src/highlight.ts | 20 ++++++++++++++++++--
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/src/features.ts b/src/features.ts
index 7bd0e98..6bdd2d6 100644
--- a/src/features.ts
+++ b/src/features.ts
@@ -44,4 +44,11 @@ export function setupAyaSpecialFeatures(context: vscode.ExtensionContext, client
     }),
     compute.applyComputedTerm,
   )));
+
+  // FIXME: workaround of https://github.com/aya-prover/aya-vscode/issues/15
+  vscode.workspace.onDidChangeTextDocument(e => {
+    const editor = vscode.window.activeTextEditor;
+    if (!editor) return;
+    hightlight.removeHighlight(editor);
+  });
 }
diff --git a/src/highlight.ts b/src/highlight.ts
index a8d863f..15ce933 100644
--- a/src/highlight.ts
+++ b/src/highlight.ts
@@ -87,12 +87,16 @@ function tokenFor(kind: Kind): Token | null {
   return null;
 }
 
+var lastHighlight: vscode.Disposable | null = null;
+var lastDecorations: Array<vscode.TextEditorDecorationType> = [];
+
 function highlightSetter(editor: vscode.TextEditor, symbol: Symbol): (target: vscode.SemanticTokensBuilder) => void {
   const color = EMACS_COLORS.get(symbol.kind);
   if (color) return (_) => {
     const decorationType = vscode.window.createTextEditorDecorationType({
       color: color,
     });
+    lastDecorations.push(decorationType);
     editor.setDecorations(decorationType, [symbol.range]);
   };
   
@@ -104,11 +108,24 @@ function highlightSetter(editor: vscode.TextEditor, symbol: Symbol): (target: vs
   };
 }
 
-var lastHighlight: vscode.Disposable | null = null;
+export function removeHighlight(editor: vscode.TextEditor) {
+  if (lastHighlight !== null) {
+    lastHighlight.dispose();
+    lastHighlight = null;
+  }
+  if (lastDecorations.length != 0) {
+    // If rangesOrOptions is empty, the existing decorations with the given decoration type will be removed.
+    // No way to remove all decorations? are you kidding me vscode?
+    lastDecorations.forEach(element => editor.setDecorations(element, []));
+    lastDecorations = [];
+  }
+}
 
 export function applyHighlight(editor: vscode.TextEditor, param: HighlightResult) {
   const selector = { language: "aya", scheme: "file" };
   const legend = new vscode.SemanticTokensLegend(TOKEN_TYPES, TOKEN_MODIFIERS);
+  // I finally find the right way to re-apply highlights.
+  removeHighlight(editor);
 
   const provider: vscode.DocumentSemanticTokensProvider = {
     provideDocumentSemanticTokens(document: vscode.TextDocument): vscode.ProviderResult<vscode.SemanticTokens> {
@@ -120,6 +137,5 @@ export function applyHighlight(editor: vscode.TextEditor, param: HighlightResult
     }
   };
 
-  if (lastHighlight !== null) lastHighlight.dispose();
   lastHighlight = vscode.languages.registerDocumentSemanticTokensProvider(selector, provider, legend);
 }