Language Server Syntax Highlights

CodeEdit/Features/Documents/CodeFileDocument/CodeFileDocument.swift

@@ -50,9 +50,6 @@ final class CodeFileDocument: NSDocument, ObservableObject {
     /// See ``CodeEditSourceEditor/CombineCoordinator``.
     @Published var contentCoordinator: CombineCoordinator = CombineCoordinator()
 
-    /// Set by ``LanguageServer`` when initialized.
-    @Published var lspCoordinator: LSPContentCoordinator?
-
     /// Used to override detected languages.
     @Published var language: CodeLanguage?
 
@@ -65,6 +62,9 @@ final class CodeFileDocument: NSDocument, ObservableObject {
     /// Document-specific overridden line wrap preference.
     @Published var wrapLines: Bool?
 
+    /// Set up by ``LanguageServer``, conforms this type to ``LanguageServerDocument``.
+    @Published var languageServerObjects: LanguageServerDocumentObjects<CodeFileDocument> = .init()
+
     /// The type of data this file document contains.
     ///
     /// If its text content is not nil, a `text` UTType is returned.
@@ -83,9 +83,6 @@ final class CodeFileDocument: NSDocument, ObservableObject {
         return type
     }
 
-    /// A stable string to use when identifying documents with language servers.
-    var languageServerURI: String? { fileURL?.absolutePath }
-
     /// Specify options for opening the file such as the initial cursor positions.
     /// Nulled by ``CodeFileView`` on first load.
     var openOptions: OpenOptions?
@@ -208,6 +205,10 @@ final class CodeFileDocument: NSDocument, ObservableObject {
         }
     }
 
+    /// Determines the code language of the document.
+    /// Use ``CodeFileDocument/language`` for the default value before using this. That property is used to override
+    /// the file's language.
+    /// - Returns: The detected code language.
     func getLanguage() -> CodeLanguage {
         guard let url = fileURL else {
             return .default
@@ -223,3 +224,13 @@ final class CodeFileDocument: NSDocument, ObservableObject {
         fileURL?.findWorkspace()
     }
 }
+
+// MARK: LanguageServerDocument
+
+extension CodeFileDocument: LanguageServerDocument {
+    /// A stable string to use when identifying documents with language servers.
+    /// Needs to be a valid URI, so always returns with the `file://` prefix to indicate it's a file URI.
+    var languageServerURI: String? {
+        fileURL?.lspURI
+    }
+}

CodeEdit/Features/Editor/Views/CodeFileView.swift

@@ -19,9 +19,13 @@ struct CodeFileView: View {
     /// The current cursor positions in the view
     @State private var cursorPositions: [CursorPosition] = []
 
+    @State private var treeSitterClient: TreeSitterClient = TreeSitterClient()
+
     /// Any coordinators passed to the view.
     private var textViewCoordinators: [TextViewCoordinator]
 
+    @State private var highlightProviders: [any HighlightProviding] = []
+
     @AppSettings(\.textEditing.defaultTabWidth)
     var defaultTabWidth
     @AppSettings(\.textEditing.indentOption)
@@ -56,16 +60,19 @@ struct CodeFileView: View {
 
     init(codeFile: CodeFileDocument, textViewCoordinators: [TextViewCoordinator] = [], isEditable: Bool = true) {
         self._codeFile = .init(wrappedValue: codeFile)
+
         self.textViewCoordinators = textViewCoordinators
             + [codeFile.contentCoordinator]
-            + [codeFile.lspCoordinator].compactMap({ $0 })
+            + [codeFile.languageServerObjects.textCoordinator].compactMap({ $0 })
         self.isEditable = isEditable
 
         if let openOptions = codeFile.openOptions {
             codeFile.openOptions = nil
             self.cursorPositions = openOptions.cursorPositions
         }
 
+        updateHighlightProviders()
+
         codeFile
             .contentCoordinator
             .textUpdatePublisher
@@ -129,6 +136,7 @@ struct CodeFileView: View {
             wrapLines: codeFile.wrapLines ?? wrapLinesToEditorWidth,
             cursorPositions: $cursorPositions,
             useThemeBackground: useThemeBackground,
+            highlightProviders: highlightProviders,
             contentInsets: edgeInsets.nsEdgeInsets,
             isEditable: isEditable,
             letterSpacing: letterSpacing,
@@ -154,6 +162,10 @@ struct CodeFileView: View {
         .onChange(of: bracketHighlight) { _ in
             bracketPairHighlight = getBracketPairHighlight()
         }
+        .onReceive(codeFile.$languageServerObjects) { languageServerObjects in
+            // This will not be called in single-file views (for now) but is safe to listen to either way
+            updateHighlightProviders(lspHighlightProvider: languageServerObjects.highlightProvider)
+        }
     }
 
     private func getBracketPairHighlight() -> BracketPairHighlight? {
@@ -174,6 +186,12 @@ struct CodeFileView: View {
             return .underline(color: color)
         }
     }
+
+    /// Updates the highlight providers array.
+    /// - Parameter lspHighlightProvider: The language server provider, if available.
+    private func updateHighlightProviders(lspHighlightProvider: HighlightProviding? = nil) {
+        highlightProviders = [lspHighlightProvider].compactMap({ $0 }) + [treeSitterClient]
+    }
 }
 
 // This extension is kept here because it should not be used elsewhere in the app and may cause confusion

CodeEdit/Features/LSP/Editor/LSPContentCoordinator.swift → CodeEdit/Features/LSP/Features/DocumentSync/LSPContentCoordinator.swift

@@ -19,7 +19,7 @@ import LanguageServerProtocol
 /// Language servers expect edits to be sent in chunks (and it helps reduce processing overhead). To do this, this class
 /// keeps an async stream around for the duration of its lifetime. The stream is sent edit notifications, which are then
 /// chunked into 250ms timed groups before being sent to the ``LanguageServer``.
-class LSPContentCoordinator: TextViewCoordinator, TextViewDelegate {
+class LSPContentCoordinator<DocumentType: LanguageServerDocument>: TextViewCoordinator, TextViewDelegate {
     // Required to avoid a large_tuple lint error
     private struct SequenceElement: Sendable {
         let uri: String
@@ -28,25 +28,27 @@ class LSPContentCoordinator: TextViewCoordinator, TextViewDelegate {
     }
 
     private var editedRange: LSPRange?
-    private var stream: AsyncStream<SequenceElement>?
     private var sequenceContinuation: AsyncStream<SequenceElement>.Continuation?
     private var task: Task<Void, Never>?
 
-    weak var languageServer: LanguageServer?
+    weak var languageServer: LanguageServer<DocumentType>?
     var documentURI: String
 
     /// Initializes a content coordinator, and begins an async stream of updates
-    init(documentURI: String, languageServer: LanguageServer) {
+    init(documentURI: String, languageServer: LanguageServer<DocumentType>) {
         self.documentURI = documentURI
         self.languageServer = languageServer
-        self.stream = AsyncStream { continuation in
-            self.sequenceContinuation = continuation
-        }
+
+        setUpUpdatesTask()
     }
 
     func setUpUpdatesTask() {
         task?.cancel()
-        guard let stream else { return }
+        // Create this stream here so it's always set up when the text view is set up, rather than only once on init.
+        let stream = AsyncStream { continuation in
+            self.sequenceContinuation = continuation
+        }
+
         task = Task.detached { [weak self] in
             // Send edit events every 250ms
             for await events in stream.chunked(by: .repeating(every: .milliseconds(250), clock: .continuous)) {

CodeEdit/Features/LSP/Features/SemanticTokens/SemanticTokenHighlightProvider.swift

@@ -0,0 +1,172 @@
+//
+//  SemanticTokenHighlightProvider.swift
+//  CodeEdit
+//
+//  Created by Khan Winter on 12/26/24.
+//
+
+import Foundation
+import LanguageServerProtocol
+import CodeEditSourceEditor
+import CodeEditTextView
+import CodeEditLanguages
+
+/// Provides semantic token information from a language server for a source editor view.
+///
+/// This class works in tangent with the ``LanguageServer`` class to ensure we don't unnecessarily request new tokens
+/// if the document isn't updated. The ``LanguageServer`` will call the
+/// ``SemanticTokenHighlightProvider/documentDidChange`` method, which in turn refreshes the semantic token storage.
+///
+/// That behavior may not be intuitive due to the
+/// ``SemanticTokenHighlightProvider/applyEdit(textView:range:delta:completion:)`` method. One might expect this class
+/// to respond to that method immediately, but it does not. It instead stores the completion passed in that method until
+/// it can respond to the edit with invalidated indices.
+final class SemanticTokenHighlightProvider<
+    Storage: GenericSemanticTokenStorage,
+    DocumentType: LanguageServerDocument
+>: HighlightProviding {
+    enum HighlightError: Error {
+        case lspRangeFailure
+    }
+
+    typealias EditCallback = @MainActor (Result<IndexSet, any Error>) -> Void
+    typealias HighlightCallback = @MainActor (Result<[HighlightRange], any Error>) -> Void
+
+    private let tokenMap: SemanticTokenMap
+    private let documentURI: String
+    private weak var languageServer: LanguageServer<DocumentType>?
+    private weak var textView: TextView?
+
+    private var lastEditCallback: EditCallback?
+    private var pendingHighlightCallbacks: [HighlightCallback] = []
+    private var storage: Storage
+
+    var documentRange: NSRange {
+        textView?.documentRange ?? .zero
+    }
+
+    init(tokenMap: SemanticTokenMap, languageServer: LanguageServer<DocumentType>, documentURI: String) {
+        self.tokenMap = tokenMap
+        self.languageServer = languageServer
+        self.documentURI = documentURI
+        self.storage = Storage()
+    }
+
+    // MARK: - Language Server Content Lifecycle
+
+    /// Called when the language server finishes sending a document update.
+    ///
+    /// This method first checks if this object has any semantic tokens. If not, requests new tokens and responds to the
+    /// `pendingHighlightCallbacks` queue with cancellation errors, causing the highlighter to re-query those indices.
+    ///
+    /// If this object already has some tokens, it determines whether or not we can request a token delta and
+    /// performs the request.
+    func documentDidChange() async throws {
+        guard let languageServer, let textView else {
+            return
+        }
+
+        guard storage.hasReceivedData else {
+            // We have no semantic token info, request it!
+            try await requestTokens(languageServer: languageServer, textView: textView)
+            await MainActor.run {
+                for callback in pendingHighlightCallbacks {
+                    callback(.failure(HighlightProvidingError.operationCancelled))
+                }
+                pendingHighlightCallbacks.removeAll()
+            }
+            return
+        }
+
+        // The document was updated. Update our token cache and send the invalidated ranges for the editor to handle.
+        if let lastResultId = storage.lastResultId {
+            try await requestDeltaTokens(languageServer: languageServer, textView: textView, lastResultId: lastResultId)
+            return
+        }
+
+        try await requestTokens(languageServer: languageServer, textView: textView)
+    }
+
+    // MARK: - LSP Token Requests
+
+    /// Requests and applies a token delta. Requires a previous response identifier.
+    private func requestDeltaTokens(
+        languageServer: LanguageServer<DocumentType>,
+        textView: TextView,
+        lastResultId: String
+    ) async throws {
+        guard let response = try await languageServer.requestSemanticTokens(
+            for: documentURI,
+            previousResultId: lastResultId
+        ) else {
+            return
+        }
+        switch response {
+        case let .optionA(tokenData):
+            await applyEntireResponse(tokenData, callback: lastEditCallback)
+        case let .optionB(deltaData):
+            await applyDeltaResponse(deltaData, callback: lastEditCallback, textView: textView)
+        }
+    }
+
+    /// Requests and applies tokens for an entire document. This does not require a previous response id, and should be
+    /// used in place of `requestDeltaTokens` when that's the case.
+    private func requestTokens(languageServer: LanguageServer<DocumentType>, textView: TextView) async throws {
+        guard let response = try await languageServer.requestSemanticTokens(for: documentURI) else {
+            return
+        }
+        await applyEntireResponse(response, callback: lastEditCallback)
+    }
+
+    // MARK: - Apply LSP Response
+
+    /// Applies a delta response from the LSP to our storage.
+    private func applyDeltaResponse(_ data: SemanticTokensDelta, callback: EditCallback?, textView: TextView?) async {
+        let lspRanges = storage.applyDelta(data)
+        lastEditCallback = nil // Don't use this callback again.
+        await MainActor.run {
+            let ranges = lspRanges.compactMap { textView?.nsRangeFrom($0) }
+            callback?(.success(IndexSet(ranges: ranges)))
+        }
+    }
+
+    private func applyEntireResponse(_ data: SemanticTokens, callback: EditCallback?) async {
+        storage.setData(data)
+        lastEditCallback = nil // Don't use this callback again.
+        await callback?(.success(IndexSet(integersIn: documentRange)))
+    }
+
+    // MARK: - Highlight Provider Conformance
+
+    func setUp(textView: TextView, codeLanguage: CodeLanguage) {
+        // Send off a request to get the initial token data
+        self.textView = textView
+        Task {
+            try await self.documentDidChange()
+        }
+    }
+
+    func applyEdit(textView: TextView, range: NSRange, delta: Int, completion: @escaping EditCallback) {
+        if let lastEditCallback {
+            lastEditCallback(.success(IndexSet())) // Don't throw a cancellation error
+        }
+        lastEditCallback = completion
+    }
+
+    func queryHighlightsFor(textView: TextView, range: NSRange, completion: @escaping HighlightCallback) {
+        guard storage.hasReceivedData else {
+            pendingHighlightCallbacks.append(completion)
+            return
+        }
+
+        guard let lspRange = textView.lspRangeFrom(nsRange: range) else {
+            completion(.failure(HighlightError.lspRangeFailure))
+            return
+        }
+        let rawTokens = storage.getTokensFor(range: lspRange)
+        let highlights = tokenMap
+            .decode(tokens: rawTokens, using: textView)
+            .filter({ $0.capture != nil || !$0.modifiers.isEmpty })
+        completion(.success(highlights))
+    }
+}

CodeEdit/Features/LSP/Editor/SemanticTokenMap.swift → CodeEdit/Features/LSP/Features/SemanticTokens/SemanticTokenMap.swift

@@ -45,20 +45,31 @@ struct SemanticTokenMap: Sendable { // swiftlint:enable line_length
     /// Decodes the compressed semantic token data into a `HighlightRange` type for use in an editor.
     /// This is marked main actor to prevent runtime errors, due to the use of the actor-isolated `rangeProvider`.
     /// - Parameters:
-    ///   - tokens: Semantic tokens from a language server.
+    ///   - tokens: Encoded semantic tokens type from a language server.
     ///   - rangeProvider: The provider to use to translate token ranges to text view ranges.
     /// - Returns: An array of decoded highlight ranges.
     @MainActor
     func decode(tokens: SemanticTokens, using rangeProvider: SemanticTokenMapRangeProvider) -> [HighlightRange] {
-        tokens.decode().compactMap { token in
+        return decode(tokens: tokens.decode(), using: rangeProvider)
+    }
+
+    /// Decodes the compressed semantic token data into a `HighlightRange` type for use in an editor.
+    /// This is marked main actor to prevent runtime errors, due to the use of the actor-isolated `rangeProvider`.
+    /// - Parameters:
+    ///   - tokens: Decoded semantic tokens from a language server.
+    ///   - rangeProvider: The provider to use to translate token ranges to text view ranges.
+    /// - Returns: An array of decoded highlight ranges.
+    @MainActor
+    func decode(tokens: [SemanticToken], using rangeProvider: SemanticTokenMapRangeProvider) -> [HighlightRange] {
+        tokens.compactMap { token in
             guard let range = rangeProvider.nsRangeFrom(line: token.line, char: token.char, length: token.length) else {
                 return nil
             }
 
+            // Only modifiers are bit packed, capture types are given as a simple index into the ``tokenTypeMap``
             let modifiers = decodeModifier(token.modifiers)
 
-            // Capture types are indicated by the index of the set bit.
-            let type = token.type > 0 ? Int(token.type.trailingZeroBitCount) : -1 // Don't try to decode 0
+            let type = Int(token.type)
             let capture = tokenTypeMap.indices.contains(type) ? tokenTypeMap[type] : nil
 
             return HighlightRange(