Load semantic tokens from a document using a sourcekitd request instead of the 0,0 edit

Sources/SourceKitD/sourcekitd_uids.swift

@@ -59,6 +59,7 @@ public struct sourcekitd_keys {
   public let request: sourcekitd_uid_t
   public let results: sourcekitd_uid_t
   public let retrieve_refactor_actions: sourcekitd_uid_t
+  public let semantic_tokens: sourcekitd_uid_t
   public let severity: sourcekitd_uid_t
   public let sourcefile: sourcekitd_uid_t
   public let sourcetext: sourcekitd_uid_t
@@ -135,6 +136,7 @@ public struct sourcekitd_keys {
     request = api.uid_get_from_cstr("key.request")!
     results = api.uid_get_from_cstr("key.results")!
     retrieve_refactor_actions = api.uid_get_from_cstr("key.retrieve_refactor_actions")!
+    semantic_tokens = api.uid_get_from_cstr("key.semantic_tokens")!
     severity = api.uid_get_from_cstr("key.severity")!
     sourcefile = api.uid_get_from_cstr("key.sourcefile")!
     sourcetext = api.uid_get_from_cstr("key.sourcetext")!
@@ -178,6 +180,7 @@ public struct sourcekitd_requests {
   public let codecomplete_close: sourcekitd_uid_t
   public let cursorinfo: sourcekitd_uid_t
   public let diagnostics: sourcekitd_uid_t
+  public let semantic_tokens: sourcekitd_uid_t
   public let expression_type: sourcekitd_uid_t
   public let find_usr: sourcekitd_uid_t
   public let variable_type: sourcekitd_uid_t
@@ -196,6 +199,7 @@ public struct sourcekitd_requests {
     codecomplete_close = api.uid_get_from_cstr("source.request.codecomplete.close")!
     cursorinfo = api.uid_get_from_cstr("source.request.cursorinfo")!
     diagnostics = api.uid_get_from_cstr("source.request.diagnostics")!
+    semantic_tokens = api.uid_get_from_cstr("source.request.semantic_tokens")!
     expression_type = api.uid_get_from_cstr("source.request.expression.type")!
     find_usr = api.uid_get_from_cstr("source.request.editor.find_usr")!
     variable_type = api.uid_get_from_cstr("source.request.variable.type")!

Sources/SourceKitLSP/CMakeLists.txt

@@ -3,7 +3,6 @@ add_library(SourceKitLSP STATIC
   CapabilityRegistry.swift
   DocumentManager.swift
   IndexStoreDB+MainFilesProvider.swift
-  RangeAdjuster.swift
   ResponseError+Init.swift
   Sequence+AsyncMap.swift
   SourceKitIndexDelegate.swift
@@ -25,7 +24,6 @@ target_sources(SourceKitLSP PRIVATE
   Swift/SemanticRefactorCommand.swift
   Swift/SemanticRefactoring.swift
   Swift/SemanticTokens.swift
-  Swift/SemanticTokensManager.swift
   Swift/SourceKitD+ResponseError.swift
   Swift/SwiftCommand.swift
   Swift/SwiftLanguageServer.swift

Sources/SourceKitLSP/CapabilityRegistry.swift

@@ -75,10 +75,6 @@ public final class CapabilityRegistry {
     clientCapabilities.workspace?.didChangeWatchedFiles?.dynamicRegistration == true
   }
 
-  public var clientHasSemanticTokenRefreshSupport: Bool {
-    clientCapabilities.workspace?.semanticTokens?.refreshSupport == true
-  }
-
   public var clientHasDiagnosticsCodeDescriptionSupport: Bool {
     clientCapabilities.textDocument?.publishDiagnostics?.codeDescriptionSupport == true
   }

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -12,11 +12,33 @@
 
 import LSPLogging
 import LanguageServerProtocol
+import SourceKitD
 import SwiftIDEUtils
 import SwiftParser
 import SwiftSyntax
 
 extension SwiftLanguageServer {
+  /// Requests the semantic highlighting tokens for the given snapshot from sourcekitd.
+  private func semanticHighlightingTokens(for snapshot: DocumentSnapshot) async throws -> [SyntaxHighlightingToken]? {
+    guard let buildSettings = await self.buildSettings(for: snapshot.uri), !buildSettings.isFallback else {
+      return nil
+    }
+
+    let skreq = SKDRequestDictionary(sourcekitd: self.sourcekitd)
+    skreq[keys.request] = requests.semantic_tokens
+    skreq[keys.sourcefile] = snapshot.uri.pseudoPath
+
+    // FIXME: SourceKit should probably cache this for us.
+    skreq[keys.compilerargs] = buildSettings.compilerArgs
+
+    let dict = try await sourcekitd.send(skreq)
+
+    guard let skTokens: SKDResponseArray = dict[keys.semantic_tokens] else {
+      return nil
+    }
+    return SyntaxHighlightingTokenParser(sourcekitd: sourcekitd).parseTokens(skTokens, in: snapshot)
+  }
+
   /// Computes an array of syntax highlighting tokens from the syntax tree that
   /// have been merged with any semantic tokens from SourceKit. If the provided
   /// range is non-empty, this function restricts its output to only those
@@ -29,13 +51,17 @@ extension SwiftLanguageServer {
     for snapshot: DocumentSnapshot,
     in range: Range<Position>? = nil
   ) async -> [SyntaxHighlightingToken] {
-    let tree = await syntaxTreeManager.syntaxTree(for: snapshot)
-    let semanticTokens = await semanticTokensManager.semanticTokens(for: snapshot.id)
+    async let tree = syntaxTreeManager.syntaxTree(for: snapshot)
+    async let semanticTokens = await orLog { try await semanticHighlightingTokens(for: snapshot) }
+
     let range =
-      range.flatMap({ $0.byteSourceRange(in: snapshot) })
-      ?? ByteSourceRange(offset: 0, length: tree.totalLength.utf8Length)
+      if let range = range.flatMap({ $0.byteSourceRange(in: snapshot) }) {
+        range
+      } else {
+        ByteSourceRange(offset: 0, length: await tree.totalLength.utf8Length)
+      }
     return
-      tree
+      await tree
       .classifications(in: range)
       .flatMap({ $0.highlightingTokens(in: snapshot) })
       .mergingTokens(with: semanticTokens ?? [])