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

This allows us to cancel the semantic tokens request. It is also the last step to allow us to open and edit documents in syntactic-only mode.

It also means that we no longer need to send a `WorkspaceSemanticTokensRefreshRequest` to the client after sourcekitd has updated its semantic tokens since with the new design the `DocumentSemanticTokensRangeRequest` will simply return the results once it has the updated semantic token.

Author: ahoppen

Sources/SourceKitD/sourcekitd_uids.swift

@@ -178,6 +178,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 +197,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/RangeAdjuster.swift

@@ -15,8 +15,27 @@ import LanguageServerProtocol
 import SwiftIDEUtils
 import SwiftParser
 import SwiftSyntax
+import SourceKitD
 
 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)
+    print(dict.description)
+    return semanticTokens(of: dict, for: 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 +48,16 @@ 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)
-    let range =
-      range.flatMap({ $0.byteSourceRange(in: snapshot) })
-      ?? ByteSourceRange(offset: 0, length: tree.totalLength.utf8Length)
+    async let tree = syntaxTreeManager.syntaxTree(for: snapshot)
+    async let semanticTokens = await orLog { try await semanticHighlightingTokens(for: snapshot) }
+    
+    let range = 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 ?? [])

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -116,7 +116,6 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
   private var inFlightPublishDiagnosticsTasks: [DocumentURI: Task<Void, Never>] = [:]
 
   let syntaxTreeManager = SyntaxTreeManager()
-  let semanticTokensManager = SemanticTokensManager()
 
   nonisolated var keys: sourcekitd_keys { return sourcekitd.keys }
   nonisolated var requests: sourcekitd_requests { return sourcekitd.requests }
@@ -195,7 +194,7 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
   }
 
   /// Returns the semantic tokens in `response` for the given `snapshot`.
-  private func semanticTokens(
+  func semanticTokens(
     of response: SKDResponseDictionary,
     for snapshot: DocumentSnapshot
   ) -> [SyntaxHighlightingToken]? {
@@ -208,48 +207,6 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
 
     return tokens
   }
-
-  /// Inform the client about changes to the syntax highlighting tokens.
-  private func requestTokensRefresh() async {
-    guard let sourceKitServer else {
-      logger.fault("Cannot request a token refresh because SourceKitServer has been destructed")
-      return
-    }
-    if capabilityRegistry.clientHasSemanticTokenRefreshSupport {
-      Task {
-        do {
-          _ = try await sourceKitServer.sendRequestToClient(WorkspaceSemanticTokensRefreshRequest())
-        } catch {
-          logger.error("refreshing tokens failed: \(error.forLogging)")
-        }
-      }
-    }
-  }
-
-  func handleDocumentUpdate(uri: DocumentURI) async {
-    guard let snapshot = documentManager.latestSnapshot(uri) else {
-      return
-    }
-    // Make the magic 0,0 replacetext request to update diagnostics and semantic tokens.
-
-    let req = SKDRequestDictionary(sourcekitd: sourcekitd)
-    req[keys.request] = requests.editor_replacetext
-    req[keys.name] = uri.pseudoPath
-    req[keys.offset] = 0
-    req[keys.length] = 0
-    req[keys.sourcetext] = ""
-
-    if let dict = try? self.sourcekitd.sendSync(req) {
-      let isSemaStage = dict[keys.diagnostic_stage] as sourcekitd_uid_t? == sourcekitd.values.diag_stage_sema
-      if isSemaStage, let semanticTokens = semanticTokens(of: dict, for: snapshot) {
-        // Only update semantic tokens if the 0,0 replacetext request returned semantic information.
-        await semanticTokensManager.setSemanticTokens(for: snapshot.id, semanticTokens: semanticTokens)
-      }
-      if isSemaStage {
-        await requestTokensRefresh()
-      }
-    }
-  }
 }
 
 extension SwiftLanguageServer {
@@ -412,8 +369,6 @@ extension SwiftLanguageServer {
     req[keys.name] = uri.pseudoPath
 
     _ = try? self.sourcekitd.sendSync(req)
-
-    await semanticTokensManager.discardSemanticTokens(for: note.textDocument.uri)
   }
 
   /// Cancels any in-flight tasks to send a `PublishedDiagnosticsNotification` after edits.
@@ -508,11 +463,6 @@ extension SwiftLanguageServer {
       postEditSnapshot: postEditSnapshot,
       edits: ConcurrentEdits(fromSequential: edits)
     )
-    await semanticTokensManager.registerEdit(
-      preEditSnapshot: preEditSnapshot.id,
-      postEditSnapshot: postEditSnapshot.id,
-      edits: note.contentChanges
-    )
 
     publishDiagnosticsIfNeeded(for: note.textDocument.uri)
   }
@@ -1343,54 +1293,12 @@ extension SwiftLanguageServer: SKDNotificationHandler {
       self.documentManager = DocumentManager()
     }
 
-    guard let dict = notification.value else {
-      logger.fault(
-        """
-        Could not decode sourcekitd notification
-        \(notification.forLogging)
-        """
-      )
-      return
-    }
-
     logger.debug(
       """
       Received notification from sourcekitd
       \(notification.forLogging)
       """
     )
-
-    if let kind: sourcekitd_uid_t = dict[self.keys.notification],
-      kind == self.values.notification_documentupdate,
-      let name: String = dict[self.keys.name]
-    {
-
-      let uri: DocumentURI
-
-      // Paths are expected to be absolute; on Windows, this means that the
-      // path is either drive letter prefixed (and thus `PathGetDriveNumberW`
-      // will provide the driver number OR it is a UNC path and `PathIsUNCW`
-      // will return `true`.  On Unix platforms, the path will start with `/`
-      // which takes care of both a regular absolute path and a POSIX
-      // alternate root path.
-
-      // TODO: this is not completely portable, e.g. MacOS 9 HFS paths are
-      // unhandled.
-      #if os(Windows)
-      let isPath: Bool = name.withCString(encodedAs: UTF16.self) {
-        !PathIsURLW($0)
-      }
-      #else
-      let isPath: Bool = name.starts(with: "/")
-      #endif
-      if isPath {
-        // If sourcekitd returns us a path, translate it back into a URL
-        uri = DocumentURI(URL(fileURLWithPath: name))
-      } else {
-        uri = DocumentURI(string: name)
-      }
-      await self.handleDocumentUpdate(uri: uri)
-    }
   }
 }