Fix #254: Force refresh semantic tokens after highlighting is generated (#271)

- Added `onDidChangeSemanticTokens` support to VS Code bindings.
- Created a global event emitter for semantic tokens in `Main`.
- Added `onUpdate` channel to `Tokens` that fires on highlighting generation.
- Subscribed to `onUpdate` in `Main` to trigger the VS Code event.
- Added regression test for `onUpdate` event emission.

Author: vic0103520

src/Editor.res

@@ -220,6 +220,20 @@ module Provider = {
   }
 
   module Mock = {
+    // https://code.visualstudio.com/api/references/vscode-api#Event
+    module Event = {
+      type t<'a>
+    }
+
+    // https://code.visualstudio.com/api/references/vscode-api#EventEmitter
+    module EventEmitter = {
+      type t<'a>
+      @module("vscode") @new external make: unit => t<'a> = "EventEmitter"
+      @get external event: t<'a> => Event.t<'a> = "event"
+      @send external fire: (t<'a>, 'a) => unit = "fire"
+      @send external dispose: t<'a> => unit = "dispose"
+    }
+
     // https://code.visualstudio.com/api/references/vscode-api#SemanticsTokens
     module SemanticsTokens = {
       type t
@@ -294,8 +308,6 @@ module Provider = {
 
     // https://code.visualstudio.com/api/references/vscode-api#DocumentSemanticTokensProvider
     module DocumentSemanticTokensProvider = {
-      // missing: onDidChangeSemanticTokens
-
       type provideDocumentSemanticTokens = (
         TextDocument.t,
         CancellationToken.t,
@@ -314,15 +326,18 @@ module Provider = {
       >
 
       type t = {
+        onDidChangeSemanticTokens: option<Event.t<unit>>,
         provideDocumentSemanticTokens: option<provideDocumentSemanticTokens>,
         provideDocumentSemanticTokensEdits: option<provideDocumentSemanticTokensEdits>,
       }
 
       let make = (
+        ~onDidChangeSemanticTokens: option<Event.t<unit>>=?,
         ~provideDocumentSemanticTokens: option<provideDocumentSemanticTokens>=?,
         ~provideDocumentSemanticTokensEdits: option<provideDocumentSemanticTokensEdits>=?,
         (),
       ) => {
+        onDidChangeSemanticTokens,
         provideDocumentSemanticTokens,
         provideDocumentSemanticTokensEdits,
       }
@@ -340,9 +355,11 @@ module Provider = {
 
   let registerDocumentSemanticTokensProvider = (
     ~provideDocumentSemanticTokens: Mock.DocumentSemanticTokensProvider.provideDocumentSemanticTokens,
+    ~onDidChangeSemanticTokens: option<Mock.Event.t<unit>>=?,
     (tokenTypes, tokenModifiers),
   ) => {
     let documentSemanticTokensProvider = Mock.DocumentSemanticTokensProvider.make(
+      ~onDidChangeSemanticTokens?,
       ~provideDocumentSemanticTokens,
       // =(textDocument, _cancel) => {
       //   let builder = Mock.SemanticTokensBuilder.makeWithLegend(semanticTokensLegend)

src/Main/Main.res

@@ -49,6 +49,7 @@ let initialize = (
   memento,
   editor,
   semanticTokensRequest,
+  semanticTokensEventEmitter,
 ) => {
   let panel = Singleton.Panel.make(extensionUri)
   // if the panel is destroyed, destroy all every State in the Registry
@@ -80,6 +81,13 @@ let initialize = (
 
   let subscribe = disposable => state.subscriptions->Array.push(disposable)->ignore
 
+  // trigger semantic tokens update when highlighting is generated
+  semanticTokensEventEmitter->Option.forEach(emitter => {
+    state.tokens->Tokens.onUpdate->Chan.on(() => {
+      emitter->Editor.Provider.Mock.EventEmitter.fire()
+    })->VSCode.Disposable.make->subscribe
+  })
+
   let getCurrentEditor = () =>
     switch VSCode.Window.activeTextEditor {
     | Some(editor) => Some(editor)
@@ -141,7 +149,7 @@ let initialize = (
   state
 }
 
-let registerDocumentSemanticTokensProvider = () => {
+let registerDocumentSemanticTokensProvider = onDidChangeSemanticTokens => {
   // these two arrays are called "legends"
   let tokenTypes = Highlighting__SemanticToken.TokenType.enumurate
   let tokenModifiers = Highlighting__SemanticToken.TokenModifier.enumurate
@@ -175,6 +183,7 @@ let registerDocumentSemanticTokensProvider = () => {
 
   Editor.Provider.registerDocumentSemanticTokensProvider(
     ~provideDocumentSemanticTokens,
+    ~onDidChangeSemanticTokens?,
     (tokenTypes, tokenModifiers),
   )
 }
@@ -269,6 +278,11 @@ let activateWithoutContext = (
   let subscribe = x => subscriptions->Array.push(x)->ignore
   let subscribeMany = xs => subscriptions->Array.pushMany(xs)->ignore
 
+  // Semantic Tokens Event
+  let semanticTokensEventEmitter = Editor.Provider.Mock.EventEmitter.make()
+  subscribe(VSCode.Disposable.make(() => semanticTokensEventEmitter->Editor.Provider.Mock.EventEmitter.dispose))
+  let onDidChangeSemanticTokens = semanticTokensEventEmitter->Editor.Provider.Mock.EventEmitter.event
+
   // Channel for testing, emits events when something has been completed,
   // for example, when the input method has translated a key sequence into a symbol
   let channels = {
@@ -363,6 +377,7 @@ let activateWithoutContext = (
           memento,
           editor,
           None,
+          Some(semanticTokensEventEmitter),
         )
         Registry.add(document, state)
       | Some(entry) =>
@@ -377,6 +392,7 @@ let activateWithoutContext = (
             memento,
             editor,
             Some(entry.semanticTokens),
+            Some(semanticTokensEventEmitter),
           )
           Registry.add(document, state)
         | Some(_) => () // should not happen
@@ -394,7 +410,7 @@ let activateWithoutContext = (
     }
   })->subscribeMany
 
-  registerDocumentSemanticTokensProvider()->subscribe
+  registerDocumentSemanticTokensProvider(Some(onDidChangeSemanticTokens))->subscribe
   registerInputMethodHintHoverProvider()->subscribe
 
   // expose the channel for testing

src/State/State.res

@@ -125,6 +125,7 @@ let destroy = async (state, alsoRemoveFromRegistry) => {
   state.channels.log->Chan.emit(Others("State.destroy: Disposed subscriptions"))
   state.channels.log->Chan.emit(TokensReset("State.destroy"))
   state.tokens->Tokens.reset
+  state.tokens->Tokens.destroyUpdateChannel
   state.channels.log->Chan.emit(Others("State.destroy: Tokens reset completed"))
   let result = await Connection.destroy(state.connection, state.channels.log)
   state.channels.log->Chan.emit(Others("State.destroy: Connection destroyed, destruction complete"))

src/Tokens.res

@@ -23,6 +23,7 @@ module type Module = {
 
   // Remove everything
   let reset: t => unit
+  let destroyUpdateChannel: t => unit
 
   // definition provider for go-to-definition
   let goToDefinition: (
@@ -41,6 +42,8 @@ module type Module = {
 
   let toOriginalOffset: (t, int) => option<int>
 
+  let onUpdate: t => Chan.t<unit>
+
   let getVSCodeTokens: t => Resource.t<array<Highlighting__SemanticToken.t>>
   let getHolePositionsFromLoad: t => Resource.t<Map.t<int, int>>
   let getHoles: t => Map.t<int, Token.t<vscodeOffset>>
@@ -118,6 +121,7 @@ module Module: Module = {
     // ranges of holes
     mutable holes: Map.t<int, Token.t<vscodeOffset>>,
     mutable holePositions: Resource.t<Map.t<int, int>>,
+    onUpdate: Chan.t<unit>,
   }
 
   let toString = self => {
@@ -153,6 +157,7 @@ module Module: Module = {
     decorations: Map.make(),
     holes: Map.make(),
     holePositions: Resource.make(),
+    onUpdate: Chan.make(),
   }
 
   let insertWithVSCodeOffsets = (self, token: Token.t<vscodeOffset>) => {
@@ -254,6 +259,8 @@ module Module: Module = {
     }
   }
 
+  let destroyUpdateChannel = self => self.onUpdate->Chan.destroy
+
   let toTokenArray = self => self.agdaTokens->AVLTree.toArray->Array.map(snd)
   let toDecorations = self => self.decorations
 
@@ -448,6 +455,8 @@ module Module: Module = {
     applyDecorations(self, editor)
     // set the holes positions
     self.holePositions->Resource.set(holePositions)
+
+    self.onUpdate->Chan.emit()
   }
 
   // Update the deltas and generate new tokens
@@ -482,6 +491,8 @@ module Module: Module = {
       None,
     )
 
+  let onUpdate = self => self.onUpdate
+
   let getVSCodeTokens = self => self.vscodeTokens
   let getHolePositionsFromLoad = self => self.holePositions
   let getHoles = self => {

test/tests/Test__Tokens.res

@@ -6,6 +6,21 @@ open Tokens
 describe("Tokens", () => {
   This.timeout(10000)
   describe("Token generation", () => {
+    Async.it(
+      "should emit `onUpdate` event when highlighting is generated",
+      async () => {
+        let ctx = await AgdaMode.makeAndLoad("GotoDefinition.agda")
+        let (promise, resolve, _) = Util.Promise_.pending()
+
+        let _disposable = ctx.state.tokens->Tokens.onUpdate->Chan.on(resolve)
+
+        ctx.state.tokens->Tokens.generateHighlighting(ctx.state.editor)
+
+        await promise
+        Assert.ok(true)
+      },
+    )
+
     Async.it(
       "should produce 28 tokens",
       async () => {