LSP Semantic Token Decoder, Improve LSP-CodeFileDoc Arch (#1951)

Author: thecoolwinter

CodeEdit.xcodeproj/project.pbxproj

@@ -33,8 +33,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/CodeEditApp/CodeEditSourceEditor",
       "state" : {
-        "revision" : "bfcde1fc536e4159ca3d596fa5b8bbbeb1524362",
-        "version" : "0.9.0"
+        "revision" : "b0688fa59fb8060840fb013afb4d6e6a96000f14",
+        "version" : "0.9.1"
       }
     },
     {

CodeEdit.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -20,7 +20,6 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
     @LazyService var lspService: LSPService
 
     func applicationDidFinishLaunching(_ notification: Notification) {
-        setupServiceContainer()
         enableWindowSizeSaveOnQuit()
         Settings.shared.preferences.general.appAppearance.applyAppearance()
         checkForFilesToOpen()
@@ -271,14 +270,6 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
             workspace.taskManager?.stopAllTasks()
         }
     }
-
-    /// Setup all the services into a ServiceContainer for the application to use.
-    @MainActor
-    private func setupServiceContainer() {
-        ServiceContainer.register(
-            LSPService()
-        )
-    }
 }
 
 extension AppDelegate {

CodeEdit/AppDelegate.swift

@@ -15,6 +15,11 @@ struct CodeEditApp: App {
     let updater: SoftwareUpdater = SoftwareUpdater()
 
     init() {
+        // Register singleton services before anything else
+        ServiceContainer.register(
+            LSPService()
+        )
+
         _ = CodeEditDocumentController.shared
         NSMenuItem.swizzle()
         NSSplitViewItem.swizzle()

CodeEdit/CodeEditApp.swift

@@ -29,7 +29,10 @@ final class CodeFileDocument: NSDocument, ObservableObject {
 
     static let logger = Logger(subsystem: Bundle.main.bundleIdentifier ?? "", category: "CodeFileDocument")
 
-    @Service var lspService: LSPService
+    /// Sent when the document is opened. The document will be sent in the notification's object.
+    static let didOpenNotification = Notification.Name(rawValue: "CodeFileDocument.didOpen")
+    /// Sent when the document is closed. The document's `fileURL` will be sent in the notification's object.
+    static let didCloseNotification = Notification.Name(rawValue: "CodeFileDocument.didClose")
 
     /// The text content of the document, stored as a text storage
     ///
@@ -47,11 +50,8 @@ final class CodeFileDocument: NSDocument, ObservableObject {
     /// See ``CodeEditSourceEditor/CombineCoordinator``.
     @Published var contentCoordinator: CombineCoordinator = CombineCoordinator()
 
-    lazy var languageServerCoordinator: LSPContentCoordinator = {
-        let coordinator = LSPContentCoordinator()
-        coordinator.uri = self.languageServerURI
-        return coordinator
-    }()
+    /// Set by ``LanguageServer`` when initialized.
+    @Published var lspCoordinator: LSPContentCoordinator?
 
     /// Used to override detected languages.
     @Published var language: CodeLanguage?
@@ -84,7 +84,7 @@ final class CodeFileDocument: NSDocument, ObservableObject {
     }
 
     /// A stable string to use when identifying documents with language servers.
-    var languageServerURI: String? { fileURL?.languageServerURI }
+    var languageServerURI: String? { fileURL?.absolutePath }
 
     /// Specify options for opening the file such as the initial cursor positions.
     /// Nulled by ``CodeFileView`` on first load.
@@ -161,6 +161,7 @@ final class CodeFileDocument: NSDocument, ObservableObject {
         } else {
             Self.logger.error("Failed to read file from data using encoding: \(rawEncoding)")
         }
+        NotificationCenter.default.post(name: Self.didOpenNotification, object: self)
     }
 
     /// Triggered when change occurred
@@ -187,7 +188,7 @@ final class CodeFileDocument: NSDocument, ObservableObject {
 
     override func close() {
         super.close()
-        lspService.closeDocument(self)
+        NotificationCenter.default.post(name: Self.didCloseNotification, object: fileURL)
     }
 
     func getLanguage() -> CodeLanguage {
@@ -202,15 +203,6 @@ final class CodeFileDocument: NSDocument, ObservableObject {
     }
 
     func findWorkspace() -> WorkspaceDocument? {
-        CodeEditDocumentController.shared.documents.first(where: { doc in
-            guard let workspace = doc as? WorkspaceDocument, let path = self.languageServerURI else { return false }
-            // createIfNotFound is safe here because it will still exit if the file and the workspace
-            // do not share a path prefix
-            return workspace
-                .workspaceFileManager?
-                .getFile(path, createIfNotFound: true)?
-                .fileDocument?
-                .isEqual(self) ?? false
-        }) as? WorkspaceDocument
+        fileURL?.findWorkspace()
     }
 }

CodeEdit/Features/Documents/CodeFileDocument/CodeFileDocument.swift

@@ -12,7 +12,7 @@ final class CodeEditDocumentController: NSDocumentController {
     @Environment(\.openWindow)
     private var openWindow
 
-    @LazyService var lspService: LSPService
+    @Service var lspService: LSPService
 
     private let fileManager = FileManager.default
 
@@ -92,13 +92,6 @@ final class CodeEditDocumentController: NSDocumentController {
             }
         }
     }
-
-    override func addDocument(_ document: NSDocument) {
-        super.addDocument(document)
-        if let document = document as? CodeFileDocument {
-            lspService.openDocument(document)
-        }
-    }
 }
 
 extension NSDocumentController {

CodeEdit/Features/Documents/Controllers/CodeEditDocumentController.swift

@@ -56,10 +56,9 @@ struct CodeFileView: View {
 
     init(codeFile: CodeFileDocument, textViewCoordinators: [TextViewCoordinator] = [], isEditable: Bool = true) {
         self._codeFile = .init(wrappedValue: codeFile)
-        self.textViewCoordinators = textViewCoordinators + [
-            codeFile.contentCoordinator,
-            codeFile.languageServerCoordinator
-        ]
+        self.textViewCoordinators = textViewCoordinators
+            + [codeFile.contentCoordinator]
+            + [codeFile.lspCoordinator].compactMap({ $0 })
         self.isEditable = isEditable
 
         if let openOptions = codeFile.openOptions {
@@ -138,7 +137,6 @@ struct CodeFileView: View {
             undoManager: undoManager,
             coordinators: textViewCoordinators
         )
-
         .id(codeFile.fileURL)
         .background {
             if colorScheme == .dark {

CodeEdit/Features/Editor/Views/CodeFileView.swift

@@ -33,9 +33,12 @@ class LSPContentCoordinator: TextViewCoordinator, TextViewDelegate {
     private var task: Task<Void, Never>?
 
     weak var languageServer: LanguageServer?
-    var uri: String?
+    var documentURI: String
 
-    init() {
+    /// Initializes a content coordinator, and begins an async stream of updates
+    init(documentURI: String, languageServer: LanguageServer) {
+        self.documentURI = documentURI
+        self.languageServer = languageServer
         self.stream = AsyncStream { continuation in
             self.sequenceContinuation = continuation
         }
@@ -71,12 +74,11 @@ class LSPContentCoordinator: TextViewCoordinator, TextViewDelegate {
     }
 
     func textView(_ textView: TextView, didReplaceContentsIn range: NSRange, with string: String) {
-        guard let uri,
-              let lspRange = editedRange else {
+        guard let lspRange = editedRange else {
             return
         }
         self.editedRange = nil
-        self.sequenceContinuation?.yield(SequenceElement(uri: uri, range: lspRange, string: string))
+        self.sequenceContinuation?.yield(SequenceElement(uri: documentURI, range: lspRange, string: string))
     }
 
     func destroy() {

CodeEdit/Features/LSP/LanguageServer/LSPContentCoordinator.swift renamed to CodeEdit/Features/LSP/Editor/LSPContentCoordinator.swift

@@ -0,0 +1,90 @@
+//
+//  SemanticTokenMap.swift
+//  CodeEdit
+//
+//  Created by Khan Winter on 11/10/24.
+//
+
+import LanguageClient
+import LanguageServerProtocol
+import CodeEditSourceEditor
+import CodeEditTextView
+
+// swiftlint:disable line_length
+/// Creates a mapping from a language server's semantic token options to a format readable by CodeEdit.
+/// Provides a convenience method for mapping tokens received from the server to highlight ranges suitable for
+/// highlighting in the editor.
+///
+/// Use this type to handle the initially received semantic highlight capabilities structures. This type will figure
+/// out how to read it into a format it can use.
+///
+/// After initialization, the map is static until the server is reinitialized. Consequently, this type is `Sendable`
+/// and immutable after initialization.
+///
+/// This type is not coupled to any text system via the use of the ``SemanticTokenMapRangeProvider``. When decoding to
+/// highlight ranges, provide a type that can provide ranges for highlighting.
+///
+/// [LSP Spec](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#semanticTokensLegend)
+struct SemanticTokenMap: Sendable { // swiftlint:enable line_length
+    private let tokenTypeMap: [CaptureName?]
+    private let modifierMap: [CaptureModifier?]
+
+    init(semanticCapability: TwoTypeOption<SemanticTokensOptions, SemanticTokensRegistrationOptions>) {
+        let legend: SemanticTokensLegend
+        switch semanticCapability {
+        case .optionA(let tokensOptions):
+            legend = tokensOptions.legend
+        case .optionB(let tokensRegistrationOptions):
+            legend = tokensRegistrationOptions.legend
+        }
+
+        tokenTypeMap = legend.tokenTypes.map { CaptureName.fromString($0) }
+        modifierMap = legend.tokenModifiers.map { CaptureModifier.fromString($0) }
+    }
+
+    /// 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.
+    ///   - 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
+            guard let range = rangeProvider.nsRangeFrom(line: token.line, char: token.char, length: token.length) else {
+                return nil
+            }
+
+            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 capture = tokenTypeMap.indices.contains(type) ? tokenTypeMap[type] : nil
+
+            return HighlightRange(
+                range: range,
+                capture: capture,
+                modifiers: modifiers
+            )
+        }
+    }
+
+    /// Decodes a raw modifier value into a set of capture modifiers.
+    /// - Parameter raw: The raw modifier integer to decode.
+    /// - Returns: A set of modifiers for highlighting.
+    func decodeModifier(_ raw: UInt32) -> CaptureModifierSet {
+        var modifiers: CaptureModifierSet = []
+        var raw = raw
+        while raw > 0 {
+            let idx = raw.trailingZeroBitCount
+            raw &= ~(1 << idx)
+            // We don't use `[safe:]` because it creates a double optional here. If someone knows how to extend
+            // a collection of optionals to make that return only a single optional this could be updated.
+            guard let modifier = modifierMap.indices.contains(idx) ? modifierMap[idx] : nil else {
+                continue
+            }
+            modifiers.insert(modifier)
+        }
+        return modifiers
+    }
+}

CodeEdit/Features/LSP/Editor/SemanticTokenMap.swift

@@ -0,0 +1,13 @@
+//
+//  SemanticTokenMapRangeProvider.swift
+//  CodeEdit
+//
+//  Created by Khan Winter on 12/19/24.
+//
+
+import Foundation
+
+@MainActor
+protocol SemanticTokenMapRangeProvider {
+    func nsRangeFrom(line: UInt32, char: UInt32, length: UInt32) -> NSRange?
+}

CodeEdit/Features/LSP/Editor/SemanticTokenMapRangeProvider.swift

@@ -19,7 +19,7 @@ extension LanguageServer {
             }
             logger.debug("Opening Document \(content.uri, privacy: .private)")
 
-            self.openFiles.addDocument(document)
+            openFiles.addDocument(document, for: self)
 
             let textDocument = TextDocumentItem(
                 uri: content.uri,
@@ -28,7 +28,8 @@ extension LanguageServer {
                 text: content.string
             )
             try await lspInstance.textDocumentDidOpen(DidOpenTextDocumentParams(textDocument: textDocument))
-            await updateIsolatedTextCoordinator(for: document)
+
+            await updateIsolatedDocument(document, coordinator: openFiles.contentCoordinator(for: document))
         } catch {
             logger.warning("addDocument: Error \(error)")
             throw error
@@ -118,10 +119,9 @@ extension LanguageServer {
         return DocumentContent(uri: uri, language: language, string: content)
     }
 
-    /// Updates the actor-isolated document's text coordinator to map to this server.
     @MainActor
-    fileprivate func updateIsolatedTextCoordinator(for document: CodeFileDocument) {
-        document.languageServerCoordinator.languageServer = self
+    private func updateIsolatedDocument(_ document: CodeFileDocument, coordinator: LSPContentCoordinator?) {
+        document.lspCoordinator = coordinator
     }
 
     // swiftlint:disable line_length

CodeEdit/Features/LSP/LanguageServer/Capabilities/LanguageServer+DocumentSync.swift

@@ -22,8 +22,14 @@ class LanguageServer {
     /// A cache to hold responses from the server, to minimize duplicate server requests
     let lspCache = LSPCache()
 
+    /// Tracks documents and their associated objects.
+    /// Use this property when adding new objects that need to track file data, or have a state associated with the
+    /// language server and a document. For example, the content coordinator.
     let openFiles: LanguageServerFileMap
 
+    /// Maps the language server's highlight config to one CodeEdit can read. See ``SemanticTokenMap``.
+    let highlightMap: SemanticTokenMap?
+
     /// The configuration options this server supports.
     var serverCapabilities: ServerCapabilities
 
@@ -49,6 +55,11 @@ class LanguageServer {
             subsystem: Bundle.main.bundleIdentifier ?? "",
             category: "LanguageServer.\(languageId.rawValue)"
         )
+        if let semanticTokensProvider = serverCapabilities.semanticTokensProvider {
+            self.highlightMap = SemanticTokenMap(semanticCapability: semanticTokensProvider)
+        } else {
+            self.highlightMap = nil // Server doesn't support semantic highlights
+        }
     }
 
     /// Creates and initializes a language server.
@@ -82,6 +93,8 @@ class LanguageServer {
         )
     }
 
+    // MARK: - Make Local Server Connection
+
     /// Creates a data channel for sending and receiving data with an LSP.
     /// - Parameters:
     ///   - languageId: The ID of the language to create the channel for.
@@ -105,6 +118,8 @@ class LanguageServer {
         }
     }
 
+    // MARK: - Get Init Params
+
     // swiftlint:disable function_body_length
     static func getInitParams(workspacePath: String) -> InitializingServer.InitializeParamsProvider {
         let provider: InitializingServer.InitializeParamsProvider = {
@@ -136,15 +151,15 @@ class LanguageServer {
                 // swiftlint:disable:next line_length
                 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#semanticTokensClientCapabilities
                 semanticTokens: SemanticTokensClientCapabilities(
-                    dynamicRegistration: true,
-                    requests: .init(range: true, delta: false),
-                    tokenTypes: [],
-                    tokenModifiers: [],
+                    dynamicRegistration: false,
+                    requests: .init(range: true, delta: true),
+                    tokenTypes: SemanticTokenTypes.allStrings,
+                    tokenModifiers: SemanticTokenModifiers.allStrings,
                     formats: [.relative],
                     overlappingTokenSupport: true,
                     multilineTokenSupport: true,
                     serverCancelSupport: true,
-                    augmentsSyntaxTokens: false
+                    augmentsSyntaxTokens: true
                 )
             )