Refactor DocumentTokens to Use the New Parser

Re-core the "lexical" token stream on top of SwiftSyntax and call into the parser to (re-)generate these "lexical tokens" when the document changes. Leave the semantic tokens alone since they take some non-zero amount of name lookup to fully resolve references to fully replicate.

Author: CodaFi

Package.swift

@@ -49,6 +49,8 @@ let package = Package(
           "SourceKitD",
           "SKSwiftPMWorkspace",
           .product(name: "SwiftToolsSupport-auto", package: "swift-tools-support-core"),
+          .product(name: "SwiftSyntax", package: "swift-syntax"),
+          .product(name: "SwiftParser", package: "swift-syntax"),
         ],
         exclude: ["CMakeLists.txt"]),
 

Sources/SourceKitLSP/DocumentManager.swift

@@ -153,17 +153,14 @@ public final class DocumentManager {
           // Remove all tokens in the updated range and shift later ones.
           let rangeAdjuster = RangeAdjuster(edit: edit)!
 
-          document.latestTokens.withMutableTokensOfEachKind { tokens in
-            tokens = Array(tokens.lazy
-              .compactMap {
-                var token = $0
-                if let adjustedRange = rangeAdjuster.adjust(token.range) {
-                  token.range = adjustedRange
-                  return token
-                } else {
-                  return nil
-                }
-              })
+          document.latestTokens.semantic = document.latestTokens.semantic.compactMap {
+            var token = $0
+            if let adjustedRange = rangeAdjuster.adjust(token.range) {
+              token.range = adjustedRange
+              return token
+            } else {
+              return nil
+            }
           }
         } else {
           // Full text replacement.

Sources/SourceKitLSP/DocumentTokens.swift

@@ -11,31 +11,92 @@
 //===----------------------------------------------------------------------===//
 
 import LanguageServerProtocol
+import SwiftSyntax
 
 /// Syntax highlighting tokens for a particular document.
 public struct DocumentTokens {
-  /// Lexical tokens, e.g. keywords, raw identifiers, ...
-  public var lexical: [SyntaxHighlightingToken] = []
+  /// The syntax tree representing the entire document.
+  public var syntaxTree: SourceFileSyntax?
   /// Semantic tokens, e.g. variable references, type references, ...
   public var semantic: [SyntaxHighlightingToken] = []
+}
 
-  private var merged: [SyntaxHighlightingToken] {
-    lexical.mergingTokens(with: semantic)
-  }
-  public var mergedAndSorted: [SyntaxHighlightingToken] {
-    merged.sorted { $0.start < $1.start }
+extension DocumentSnapshot {
+  private func highlightingToken(from syntax: SyntaxClassifiedRange) -> [SyntaxHighlightingToken] {
+    guard let (kind, modifiers) = syntax.kind.parseKindAndModifiers else {
+      return []
+    }
+
+    guard
+      let start: Position = self.positionOf(utf8Offset: syntax.offset),
+      let end: Position = self.positionOf(utf8Offset: syntax.endOffset)
+    else {
+      return []
+    }
+
+    let multiLineRange = start..<end
+    let ranges = multiLineRange.splitToSingleLineRanges(in: self)
+
+    return ranges.map {
+      SyntaxHighlightingToken(
+        range: $0,
+        kind: kind,
+        modifiers: modifiers
+      )
+    }
   }
 
-  /// Modifies the syntax highlighting tokens of each kind
-  /// (lexical and semantic) according to `action`.
-  public mutating func withMutableTokensOfEachKind(_ action: (inout [SyntaxHighlightingToken]) -> Void) {
-    action(&lexical)
-    action(&semantic)
+  /// 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
+  /// tokens whose ranges overlap it. If no range is provided, tokens for the
+  /// entire document are returned.
+  ///
+  /// - Parameter range: The range of tokens to restrict this function to, if any.
+  /// - Returns: An array of syntax highlighting tokens.
+  public func mergedAndSortedTokens(in range: Range<Position>? = nil) -> [SyntaxHighlightingToken] {
+    guard let tree = self.tokens.syntaxTree else {
+      return self.tokens.semantic
+    }
+    let range = range.flatMap({ self.utf8OffsetRange(of: $0) }).map({ ByteSourceRange(offset: $0.startIndex, length: $0.count) })
+             ?? ByteSourceRange(offset: 0, length: tree.byteSize)
+    return tree
+      .classifications(in: range)
+      .flatMap({ self.highlightingToken(from: $0) })
+      .mergingTokens(with: self.tokens.semantic)
+      .sorted { $0.start < $1.start }
   }
+}
 
-  // Replace all lexical tokens in `range`.
-  public mutating func replaceLexical(in range: Range<Position>, with newTokens: [SyntaxHighlightingToken]) {
-    lexical.removeAll { $0.range.overlaps(range) }
-    lexical += newTokens
+extension SyntaxClassification {
+  fileprivate var parseKindAndModifiers: (SyntaxHighlightingToken.Kind, SyntaxHighlightingToken.Modifiers)? {
+    switch self {
+    case .none:
+      return nil
+    case .editorPlaceholder:
+      return nil
+    case .stringInterpolationAnchor:
+      return nil
+    case .keyword:
+      return (.keyword, [])
+    case .identifier, .typeIdentifier, .dollarIdentifier:
+      return (.identifier, [])
+    case .operatorIdentifier:
+      return (.operator, [])
+    case .integerLiteral, .floatingLiteral:
+      return (.number, [])
+    case .stringLiteral:
+      return (.string, [])
+    case .poundDirectiveKeyword:
+      return (.macro, [])
+    case .buildConfigId, .objectLiteral:
+      return (.macro, [])
+    case .attribute:
+      return (.modifier, [])
+    case .lineComment, .blockComment:
+      return (.comment, [])
+    case .docLineComment, .docBlockComment:
+      return (.comment, .documentation)
+    }
   }
 }

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -17,6 +17,8 @@ import LSPLogging
 import SKCore
 import SKSupport
 import SourceKitD
+import SwiftSyntax
+import SwiftParser
 
 #if os(Windows)
 import WinSDK
@@ -168,13 +170,12 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
   /// Updates the lexical tokens for the given `snapshot`.
   /// Must be called on `self.queue`.
   private func updateLexicalTokens(
-    response: SKDResponseDictionary,
     for snapshot: DocumentSnapshot
   ) {
     dispatchPrecondition(condition: .onQueue(queue))
 
     let uri = snapshot.document.uri
-    let docTokens = updatedLexicalTokens(response: response, for: snapshot)
+    let docTokens = updateSyntaxTree(for: snapshot)
 
     do {
       try documentManager.updateTokens(uri, tokens: docTokens)
@@ -184,31 +185,13 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
   }
 
   /// Returns the updated lexical tokens for the given `snapshot`.
-  private func updatedLexicalTokens(
-    response: SKDResponseDictionary,
+  private func updateSyntaxTree(
     for snapshot: DocumentSnapshot
   ) -> DocumentTokens {
     logExecutionTime(level: .debug) {
       var docTokens = snapshot.tokens
 
-      guard let offset: Int = response[keys.offset],
-            let length: Int = response[keys.length],
-            let start: Position = snapshot.positionOf(utf8Offset: offset),
-            let end: Position = snapshot.positionOf(utf8Offset: offset + length) else {
-        // This e.g. happens in the case of empty edits
-        log("did not update lexical/syntactic tokens, no range found", level: .debug)
-        return docTokens
-      }
-
-      let range = start..<end
-
-      if let syntaxMap: SKDResponseArray = response[keys.syntaxmap] {
-        let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd)
-        var tokens: [SyntaxHighlightingToken] = []
-        tokenParser.parseTokens(syntaxMap, in: snapshot, into: &tokens)
-
-        docTokens.replaceLexical(in: range, with: tokens)
-      }
+      docTokens.syntaxTree = Parser.parse(source: snapshot.text)
 
       return docTokens
     }
@@ -438,7 +421,7 @@ extension SwiftLanguageServer {
     }
     self.publishDiagnostics(
         response: dict, for: snapshot, compileCommand: compileCmd)
-    self.updateLexicalTokens(response: dict, for: snapshot)
+    self.updateLexicalTokens(for: snapshot)
   }
 
   public func documentUpdatedBuildSettings(_ uri: DocumentURI, change: FileBuildSettingsChange) {
@@ -503,7 +486,7 @@ extension SwiftLanguageServer {
         return
       }
       self.publishDiagnostics(response: dict, for: snapshot, compileCommand: compileCommand)
-      self.updateLexicalTokens(response: dict, for: snapshot)
+      self.updateLexicalTokens(for: snapshot)
     }
   }
 
@@ -557,8 +540,8 @@ extension SwiftLanguageServer {
 
         self.adjustDiagnosticRanges(of: note.textDocument.uri, for: edit)
       } updateDocumentTokens: { (after: DocumentSnapshot) in
-        if let dict = lastResponse {
-          return self.updatedLexicalTokens(response: dict, for: after)
+        if lastResponse != nil {
+          return self.updateSyntaxTree(for: after)
         } else {
           return DocumentTokens()
         }
@@ -876,7 +859,7 @@ extension SwiftLanguageServer {
         return
       }
 
-      let tokens = snapshot.tokens.mergedAndSorted
+      let tokens = snapshot.mergedAndSortedTokens()
       let encodedTokens = tokens.lspEncoded
 
       req.reply(DocumentSemanticTokensResponse(data: encodedTokens))
@@ -899,7 +882,7 @@ extension SwiftLanguageServer {
         return
       }
 
-      let tokens = snapshot.tokens.mergedAndSorted.filter { $0.range.overlaps(range) }
+      let tokens = snapshot.mergedAndSortedTokens(in: range)
       let encodedTokens = tokens.lspEncoded
 
       req.reply(DocumentSemanticTokensResponse(data: encodedTokens))

Sources/SourceKitLSP/Swift/SyntaxHighlightingTokenParser.swift

@@ -187,7 +187,7 @@ struct SyntaxHighlightingTokenParser {
 
 extension Range where Bound == Position {
   /// Splits a potentially multi-line range to multiple single-line ranges.
-  fileprivate func splitToSingleLineRanges(in snapshot: DocumentSnapshot) -> [Self] {
+  func splitToSingleLineRanges(in snapshot: DocumentSnapshot) -> [Self] {
     if isEmpty {
       return []
     }