feat(highlight): Support for semanticTokens/full/delta

CHANGELOG.md

@@ -9,6 +9,9 @@
 #### LSP
 - Temporary "increase timeout" toggle (currently, it bumps all timeouts to 15 seconds)
 
+#### Highlighting
+- Support for Semantic Token Deltas (more compact way for the LSP server to send back data when typing fast)
+
 ### Fixed
 
 #### Highlighting

src/main/java/com/falsepattern/zigbrains/HighlightingUtil.java

@@ -18,13 +18,23 @@ package com.falsepattern.zigbrains;
 
 import com.falsepattern.zigbrains.lsp.ZLSEditorEventManager;
 import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.editor.RangeMarker;
 import com.intellij.openapi.editor.markup.HighlighterLayer;
 import com.intellij.openapi.editor.markup.HighlighterTargetArea;
+import com.intellij.openapi.editor.markup.RangeHighlighter;
+import com.intellij.openapi.util.Key;
 import org.wso2.lsp4intellij.editor.EditorEventManager;
 
-import javax.swing.SwingUtilities;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Stream;
+
 
 public class HighlightingUtil {
+    private static final Key<Integer> HL_HASH = Key.create("HIGHLIGHTING_HASH");
     public static void refreshHighlighting(EditorEventManager eem) {
         var app = ApplicationManager.getApplication();
         if (!(eem instanceof ZLSEditorEventManager manager)) {
@@ -34,16 +44,55 @@ public class HighlightingUtil {
         if (editor == null) {
             return;
         }
-        app.runReadAction(() -> {
+        app.executeOnPooledThread(() -> {
-            var highlightRanges = manager.semanticHighlighting();
+            if (editor.isDisposed()) {
-            var markup = editor.getMarkupModel();
+                return;
-            SwingUtilities.invokeLater(() -> ApplicationManager.getApplication().runWriteAction(() -> {
-                markup.removeAllHighlighters();
-                for (var range : highlightRanges) {
-                    markup.addRangeHighlighter(range.color(), range.start(), range.end(), HighlighterLayer.SYNTAX,
-                                               HighlighterTargetArea.EXACT_RANGE);
             }
-            }));
+            var highlightRanges = manager.semanticHighlighting();
+            var newHash = highlightRanges.hashCode();
+            app.invokeAndWait(() -> {
+                if (editor.isDisposed()) {
+                    return;
+                }
+                var markup = editor.getMarkupModel();
+                var hash = markup.getUserData(HL_HASH);
+                if (hash != null && hash == newHash) {
+                    return;
+                }
+                markup.putUserData(HL_HASH, newHash);
+                List<Map.Entry<RangeHighlighter, Integer>> highlightersSorted = null;
+                var documentLength = editor.getDocument().getTextLength();
+                for (var range : highlightRanges) {
+                    if (range.start() == 0 && range.remove() == -1 && highlightRanges.size() == 1) {
+                        markup.removeAllHighlighters();
+                    } else if (highlightersSorted == null) {
+                        highlightersSorted = new ArrayList<>(Stream.of(markup.getAllHighlighters())
+                                                                   .map(hl -> Map.entry(hl, hl.getStartOffset()))
+                                                                   .sorted(Comparator.comparingInt(Map.Entry::getValue))
+                                                                   .toList());
+                    }
+                    var start = range.start();
+                    var toRemove = range.remove();
+                    if (toRemove > 0) {
+                        for (int i = 0; i < highlightersSorted.size(); i++) {
+                            var hl = highlightersSorted.get(i);
+                            if (hl.getValue() >= start) {
+                                for (int j = 0; j < toRemove; j++) {
+                                    highlightersSorted.remove(i);
+                                    markup.removeHighlighter(hl.getKey());
+                                }
+                            }
+                        }
+                    }
+                    for (var edit: range.add()) {
+                        var end = edit.end();
+                        if (end > documentLength - 1) {
+                            end = documentLength - 1;
+                        }
+                        markup.addRangeHighlighter(edit.color(), edit.start(), end, HighlighterLayer.SYNTAX, HighlighterTargetArea.EXACT_RANGE);
+                    }
+                }
+            });
         });
     }
 }

src/main/java/com/falsepattern/zigbrains/ide/SemaEdit.java

@@ -0,0 +1,6 @@
+package com.falsepattern.zigbrains.ide;
+
+import java.util.List;
+
+public record SemaEdit(int start, int remove, List<SemaRange> add) {
+}

src/main/java/com/falsepattern/zigbrains/lsp/ZLSEditorEventManager.java

@@ -16,15 +16,20 @@
 
 package com.falsepattern.zigbrains.lsp;
 
-import com.falsepattern.zigbrains.ide.SemaRange;
+import com.falsepattern.zigbrains.ide.SemaEdit;
 import com.falsepattern.zigbrains.util.TokenDecoder;
 import com.intellij.lang.annotation.Annotation;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.editor.event.DocumentListener;
 import com.intellij.openapi.editor.event.EditorMouseListener;
 import com.intellij.openapi.editor.event.EditorMouseMotionListener;
+import org.eclipse.lsp4j.SemanticTokens;
+import org.eclipse.lsp4j.SemanticTokensDelta;
+import org.eclipse.lsp4j.SemanticTokensDeltaParams;
+import org.eclipse.lsp4j.SemanticTokensEdit;
 import org.eclipse.lsp4j.SemanticTokensParams;
 import org.eclipse.lsp4j.jsonrpc.JsonRpcException;
+import org.eclipse.lsp4j.jsonrpc.messages.Either;
 import org.wso2.lsp4intellij.client.languageserver.ServerOptions;
 import org.wso2.lsp4intellij.client.languageserver.requestmanager.RequestManager;
 import org.wso2.lsp4intellij.client.languageserver.wrapper.LanguageServerWrapper;
@@ -34,6 +39,7 @@ import org.wso2.lsp4intellij.requests.Timeouts;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -41,6 +47,7 @@ import java.util.concurrent.TimeoutException;
 import static org.wso2.lsp4intellij.requests.Timeout.getTimeout;
 
 public class ZLSEditorEventManager extends EditorEventManager {
+    private static String previousResultID = null;
     public ZLSEditorEventManager(Editor editor, DocumentListener documentListener, EditorMouseListener mouseListener, EditorMouseMotionListener mouseMotionListener, LSPCaretListenerImpl caretListener, RequestManager requestmanager, ServerOptions serverOptions, LanguageServerWrapper wrapper) {
         super(editor, documentListener, mouseListener, mouseMotionListener, caretListener, requestmanager,
               serverOptions, wrapper);
@@ -51,8 +58,8 @@ public class ZLSEditorEventManager extends EditorEventManager {
         return super.getAnnotations();
     }
 
-    public List<SemaRange> semanticHighlighting() {
+    public List<SemaEdit> semanticHighlighting() {
-        var result = new ArrayList<SemaRange>();
+        var result = new ArrayList<SemaEdit>();
         if (!(getRequestManager() instanceof ZLSRequestManager requestManager)) {
             return result;
         }
@@ -61,9 +68,14 @@ public class ZLSEditorEventManager extends EditorEventManager {
             return result;
         }
         var legend = legendOptional.get();
-        var request = requestManager.semanticTokens(new SemanticTokensParams(getIdentifier()));
+        CompletableFuture<Either<SemanticTokens, SemanticTokensDelta>> request = null;
-        if (request == null) {
+        if (previousResultID == null) {
-            return result;
+            var param = new SemanticTokensParams(getIdentifier());
+            request = requestManager.semanticTokensFull(param)
+                                    .thenApply(tokens -> tokens != null ? Either.forLeft(tokens) : null);
+        } else {
+            var param = new SemanticTokensDeltaParams(getIdentifier(), previousResultID);
+            request = requestManager.semanticTokensFullDelta(param);
         }
 
         try {
@@ -72,8 +84,20 @@ public class ZLSEditorEventManager extends EditorEventManager {
             if (res == null) {
                 return result;
             }
-            var responseData = res.getData();
+            if (res.isLeft()) {
-            return TokenDecoder.decodePayload(editor, legend, responseData);
+                var response = res.getLeft();
+                previousResultID = response.getResultId();
+                var responseData = response.getData();
+                result.add(new SemaEdit(0, -1, TokenDecoder.decodePayload(0, editor, legend, responseData)));
+            } else {
+                var response = res.getRight();
+                previousResultID = response.getResultId();
+                var edits = response.getEdits();
+                for (SemanticTokensEdit edit : edits) {
+                    var add = TokenDecoder.decodePayload(0, editor, legend, edit.getData());
+                    result.add(new SemaEdit(edit.getStart(), edit.getDeleteCount(), add));
+                }
+            }
         } catch (TimeoutException | InterruptedException e) {
             LOG.warn(e);
             wrapper.notifyFailure(Timeouts.COMPLETION);

src/main/java/com/falsepattern/zigbrains/lsp/ZLSRequestManager.java

@@ -20,16 +20,23 @@ import com.falsepattern.zigbrains.HighlightingUtil;
 import com.intellij.openapi.diagnostic.Logger;
 import org.eclipse.lsp4j.DidChangeTextDocumentParams;
 import org.eclipse.lsp4j.DidOpenTextDocumentParams;
+import org.eclipse.lsp4j.FoldingRange;
+import org.eclipse.lsp4j.FoldingRangeRequestParams;
 import org.eclipse.lsp4j.SemanticTokens;
+import org.eclipse.lsp4j.SemanticTokensDelta;
+import org.eclipse.lsp4j.SemanticTokensDeltaParams;
 import org.eclipse.lsp4j.SemanticTokensLegend;
 import org.eclipse.lsp4j.SemanticTokensParams;
 import org.eclipse.lsp4j.SemanticTokensWithRegistrationOptions;
 import org.eclipse.lsp4j.ServerCapabilities;
+import org.eclipse.lsp4j.jsonrpc.messages.Either;
 import org.eclipse.lsp4j.services.LanguageClient;
 import org.eclipse.lsp4j.services.LanguageServer;
 import org.wso2.lsp4intellij.client.languageserver.requestmanager.DefaultRequestManager;
 import org.wso2.lsp4intellij.client.languageserver.wrapper.LanguageServerWrapper;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
@@ -40,7 +47,7 @@ public class ZLSRequestManager extends DefaultRequestManager {
         super(wrapper, server, client, serverCapabilities);
     }
 
-    public CompletableFuture<SemanticTokens> semanticTokens(SemanticTokensParams params) {
+    public CompletableFuture<SemanticTokens> semanticTokensFull(SemanticTokensParams params) {
         if (checkStatus()) {
             try {
                 return (getServerCapabilities().getSemanticTokensProvider() != null)
@@ -53,6 +60,19 @@ public class ZLSRequestManager extends DefaultRequestManager {
         return null;
     }
 
+    public CompletableFuture<Either<SemanticTokens, SemanticTokensDelta>> semanticTokensFullDelta(SemanticTokensDeltaParams params) {
+        if (checkStatus()) {
+            try {
+                return (getServerCapabilities().getSemanticTokensProvider() != null)
+                       ? getTextDocumentService().semanticTokensFullDelta(params) : null;
+            } catch (Exception e) {
+                crashed(e);
+                return null;
+            }
+        }
+        return null;
+    }
+
     @Override
     public void didChange(DidChangeTextDocumentParams params) {
         super.didChange(params);

src/main/java/com/falsepattern/zigbrains/util/TokenDecoder.java

@@ -44,12 +44,12 @@ public class TokenDecoder {
         }
     }
 
-    public static List<SemaRange> decodePayload(Editor editor, SemanticTokensLegend legend, List<Integer> responseData) {
+    public static List<SemaRange> decodePayload(int baseOffset, Editor editor, SemanticTokensLegend legend, List<Integer> responseData) {
         var result = new ArrayList<SemaRange>();
         var application = ApplicationManager.getApplication();
         int dataSize = responseData.size();
-
+        var startPos = application.runReadAction((Computable<LogicalPosition>)() -> editor.offsetToLogicalPosition(baseOffset));
-        Token prevToken = null;
+        Token prevToken = new Token(startPos.line, startPos.column, 0, 0, 0);
         var types = legend.getTokenTypes();
         var modifiers = legend.getTokenModifiers();
         var modCount = Math.min(31, modifiers.size());