Apply stylistic suggestions regarding semantic tokens

- Rename SemanticToken to SyntaxHighlightingToken
  Since we use these tokens for more than just semantic highlighting and
  they are an internal structure, we rename them for clarity.
- Move support methods into Array<SyntaxHighlightingToken>
- Move syntax highlighting token decoder into SKTestSupport
  The decoder is only needed for testing.
- Improve some doc comments in SyntaxHighlightingToken
- Update SyntaxHighlightingToken.Kind to match the LSP spec
- Update SyntaxHighlightingToken.Modifiers.lspName
- Clean up LSP encoding of tokens
- Rename DocumentTokens.sorted to .mergedAndSorted
- Slightly improve naming in addLexicalTokens
- Rename DocumenTokens.withEachKind
- Add doc comment clarifying DocumentSnapshot.tokens
- Pass snapshot to parse methods in token parser
- Update function/method classifications for tokens
- Log unknown token kinds
- Require token updates to be on SwiftLanguageServer's queue
- Move token refresh request into new method

Author: fwcd

Sources/SKTestSupport/Array+SyntaxHighlightingToken.swift

@@ -0,0 +1,51 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SourceKitLSP
+import LanguageServerProtocol
+
+extension Array where Element == SyntaxHighlightingToken {
+  /// Decodes the LSP representation of syntax highlighting tokens
+  public init(lspEncodedTokens rawTokens: [UInt32]) {
+    self.init()
+    assert(rawTokens.count.isMultiple(of: 5))
+    reserveCapacity(rawTokens.count / 5)
+
+    var current = Position(line: 0, utf16index: 0)
+
+    for i in stride(from: 0, to: rawTokens.count, by: 5) {
+      let lineDelta = Int(rawTokens[i])
+      let charDelta = Int(rawTokens[i + 1])
+      let length = Int(rawTokens[i + 2])
+      let rawKind = rawTokens[i + 3]
+      let rawModifiers = rawTokens[i + 4]
+
+      guard let kind = SyntaxHighlightingToken.Kind(rawValue: rawKind) else { continue }
+      let modifiers = SyntaxHighlightingToken.Modifiers(rawValue: rawModifiers)
+
+      current.line += lineDelta
+
+      if lineDelta == 0 {
+        current.utf16index += charDelta
+      } else {
+        current.utf16index = charDelta
+      }
+
+      append(SyntaxHighlightingToken(
+        start: current,
+        length: length,
+        kind: kind,
+        modifiers: modifiers
+      ))
+    }
+  }
+}

Sources/SourceKitLSP/CMakeLists.txt

@@ -27,10 +27,10 @@ target_sources(SourceKitLSP PRIVATE
   Swift/ExpressionTypeInfo.swift
   Swift/SemanticRefactorCommand.swift
   Swift/SemanticRefactoring.swift
-  Swift/SemanticTokenParser.swift
   Swift/SourceKitD+ResponseError.swift
   Swift/SwiftCommand.swift
   Swift/SwiftLanguageServer.swift
+  Swift/SyntaxHighlightingToken.swift
   Swift/VariableTypeInfo.swift)
 set_target_properties(SourceKitLSP PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})

Sources/SourceKitLSP/DocumentManager.swift

@@ -17,20 +17,22 @@ import SKSupport
 
 public struct DocumentTokens {
   /// Lexical tokens, e.g. keywords, raw identifiers, ...
-  public var lexical: [SemanticToken] = []
+  public var lexical: [SyntaxHighlightingToken] = []
   /// Syntactic tokens, e.g. declarations, etc.
-  public var syntactic: [SemanticToken] = []
+  public var syntactic: [SyntaxHighlightingToken] = []
   /// Semantic tokens, e.g. variable references, type references, ...
-  public var semantic: [SemanticToken] = []
+  public var semantic: [SyntaxHighlightingToken] = []
 
-  public var merged: [SemanticToken] {
-    [lexical, syntactic, semantic].reduce([], mergeSemanticTokens)
+  public var merged: [SyntaxHighlightingToken] {
+    [lexical, syntactic, semantic].reduce([]) { $0.mergingTokens(with: $1) }
   }
-  public var sorted: [SemanticToken] {
+  public var mergedAndSorted: [SyntaxHighlightingToken] {
     merged.sorted { $0.start < $1.start }
   }
 
-  public mutating func withEachKind(_ action: (inout [SemanticToken]) -> Void) {
+  /// Modifies the syntax highlighting tokens of each kind
+  /// (lexical, syntactic, semantic) according to `action`.
+  public mutating func withMutableTokensOfEachKind(_ action: (inout [SyntaxHighlightingToken]) -> Void) {
     action(&lexical)
     action(&syntactic)
     action(&semantic)
@@ -41,6 +43,9 @@ public struct DocumentSnapshot {
   public var document: Document
   public var version: Int
   public var lineTable: LineTable
+  /// Syntax highlighting tokens for the document. Note that
+  /// `uri` + `latestVersion` only uniquely identifies a snapshot's content,
+  /// the tokens are updated independently and only used internally.
   public var tokens: DocumentTokens
 
   public var text: String { lineTable.content }
@@ -173,7 +178,7 @@ public final class DocumentManager {
             character.isWhitespace || character.isPunctuation || character.isSymbol
           }
 
-          func update(tokens: inout [SemanticToken]) {
+          func update(tokens: inout [SyntaxHighlightingToken]) {
             tokens = Array(tokens.lazy
               .filter {
                 // Only keep tokens that don't overlap with or are directly
@@ -197,7 +202,7 @@ public final class DocumentManager {
               })
           }
 
-          document.latestTokens.withEachKind(update(tokens:))
+          document.latestTokens.withMutableTokensOfEachKind(update(tokens:))
         } else {
           // Full text replacement.
           document.latestLineTable = LineTable(edit.text)
@@ -218,7 +223,7 @@ public final class DocumentManager {
   @discardableResult
   public func replaceSemanticTokens(
     _ uri: DocumentURI,
-    tokens: [SemanticToken]
+    tokens: [SyntaxHighlightingToken]
   ) throws -> DocumentSnapshot {
     return try queue.sync {
       guard let document = documents[uri] else {
@@ -237,7 +242,7 @@ public final class DocumentManager {
   @discardableResult
   public func replaceSyntacticTokens(
     _ uri: DocumentURI,
-    tokens: [SemanticToken]
+    tokens: [SyntaxHighlightingToken]
   ) throws -> DocumentSnapshot {
     return try queue.sync {
       guard let document = documents[uri] else {
@@ -256,24 +261,24 @@ public final class DocumentManager {
   @discardableResult
   public func addLexicalTokens(
     _ uri: DocumentURI,
-    tokens: [SemanticToken]
+    tokens newTokens: [SyntaxHighlightingToken]
   ) throws -> DocumentSnapshot {
     return try queue.sync {
       guard let document = documents[uri] else {
         throw Error.missingDocument(uri)
       }
 
-      if !tokens.isEmpty {
+      if !newTokens.isEmpty {
         // Remove all tokens that overlap with previous tokens
 
-        func removeAllOverlapping(tokens existingTokens: inout [SemanticToken]) {
+        func removeAllOverlapping(tokens existingTokens: inout [SyntaxHighlightingToken]) {
           existingTokens.removeAll { existing in
-            tokens.contains { existing.range.overlaps($0.range) }
+            newTokens.contains { existing.range.overlaps($0.range) }
           }
         }
 
-        document.latestTokens.withEachKind(removeAllOverlapping(tokens:))
-        document.latestTokens.lexical += tokens
+        document.latestTokens.withMutableTokensOfEachKind(removeAllOverlapping(tokens:))
+        document.latestTokens.lexical += newTokens
       }
 
       return document.latestSnapshot

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -162,18 +162,18 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
   }
 
   /// Update lexical and syntactic tokens for the given `snapshot`.
-  func updateLexicalAndSyntacticTokens(
+  /// Should only be called on `queue`.
+  private func updateLexicalAndSyntacticTokens(
     response: SKDResponseDictionary,
     for snapshot: DocumentSnapshot
   ) {
+    dispatchPrecondition(condition: .onQueue(queue))
+
     let uri = snapshot.document.uri
 
     if let syntaxMap: SKDResponseArray = response[keys.syntaxmap] {
-      let tokenParser = SemanticTokenParser(
-        sourcekitd: sourcekitd,
-        snapshot: snapshot
-      )
-      let tokens = tokenParser.parseTokens(syntaxMap)
+      let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd)
+      let tokens = tokenParser.parseTokens(syntaxMap, in: snapshot)
       do {
         try documentManager.addLexicalTokens(uri, tokens: tokens)
       } catch {
@@ -182,50 +182,51 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
     }
 
     if let substructure: SKDResponseArray = response[keys.substructure] {
-      let tokenParser = SemanticTokenParser(
-        sourcekitd: sourcekitd,
-        snapshot: snapshot,
-        useName: true
-      )
-      let tokens = tokenParser.parseTokens(substructure)
+      let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd, useName: true)
+      let tokens = tokenParser.parseTokens(substructure, in: snapshot)
       do {
         try documentManager.replaceSyntacticTokens(uri, tokens: tokens)
       } catch {
-        log("updating lexical tokens for \(uri) failed: \(error)", level: .warning)
+        log("updating syntactic tokens for \(uri) failed: \(error)", level: .warning)
       }
     }
   }
 
   /// Update semantic tokens for the given `snapshot`.
-  func updateSemanticTokens(
+  /// Should only be called on `queue`.
+  private func updateSemanticTokens(
     response: SKDResponseDictionary,
     for snapshot: DocumentSnapshot
   ) {
+    dispatchPrecondition(condition: .onQueue(queue))
+
     let uri = snapshot.document.uri
     guard let skTokens: SKDResponseArray = response[keys.annotations] else {
       return
     }
 
-    let tokenParser = SemanticTokenParser(
-      sourcekitd: sourcekitd,
-      snapshot: snapshot
-    )
-    let tokens = tokenParser.parseTokens(skTokens)
+    let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd)
+    let tokens = tokenParser.parseTokens(skTokens, in: snapshot)
 
     do {
       try documentManager.replaceSemanticTokens(uri, tokens: tokens)
-      if clientCapabilities.workspace?.semanticTokens?.refreshSupport ?? false {
-        _ = client.send(WorkspaceSemanticTokensRefreshRequest(), queue: queue) { result in
-          if let error = result.failure {
-            log("refreshing semantic tokens for \(uri) failed: \(error)", level: .warning)
-          }
-        }
-      }
+      
     } catch {
       log("updating semantic tokens for \(uri) failed: \(error)", level: .warning)
     }
   }
 
+  /// Inform the client about changes to the syntax highlighting tokens.
+  private func requestTokensRefresh() {
+    if clientCapabilities.workspace?.semanticTokens?.refreshSupport ?? false {
+      _ = client.send(WorkspaceSemanticTokensRefreshRequest(), queue: queue) { result in
+        if let error = result.failure {
+          log("refreshing tokens failed: \(error)", level: .warning)
+        }
+      }
+    }
+  }
+
   /// Publish diagnostics for the given `snapshot`. We withhold semantic diagnostics if we are using
   /// fallback arguments.
   ///
@@ -288,6 +289,7 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
     if let dict = try? self.sourcekitd.sendSync(req) {
       publishDiagnostics(response: dict, for: snapshot, compileCommand: compileCommand)
       updateSemanticTokens(response: dict, for: snapshot)
+      requestTokensRefresh()
     }
   }
 }
@@ -323,8 +325,8 @@ extension SwiftLanguageServer {
         commands: builtinSwiftCommands),
       semanticTokensProvider: SemanticTokensOptions(
         legend: SemanticTokensLegend(
-          tokenTypes: SemanticToken.Kind.allCases.map(\.lspName),
-          tokenModifiers: SemanticToken.Modifiers.allCases.compactMap(\.lspName)),
+          tokenTypes: SyntaxHighlightingToken.Kind.allCases.map(\.lspName),
+          tokenModifiers: SyntaxHighlightingToken.Modifiers.allCases.map { $0.lspName! }),
         range: .bool(true),
         full: .bool(true))
     ))
@@ -788,8 +790,8 @@ extension SwiftLanguageServer {
         return
       }
 
-      let tokens = snapshot.tokens.sorted
-      let encodedTokens = encodeToIntArray(semanticTokens: tokens)
+      let tokens = snapshot.tokens.mergedAndSorted
+      let encodedTokens = tokens.lspEncoded
 
       req.reply(DocumentSemanticTokensResponse(data: encodedTokens))
     }
@@ -811,8 +813,8 @@ extension SwiftLanguageServer {
         return
       }
 
-      let tokens = snapshot.tokens.sorted.filter { $0.range.overlaps(range) }
-      let encodedTokens = encodeToIntArray(semanticTokens: tokens)
+      let tokens = snapshot.tokens.mergedAndSorted.filter { $0.range.overlaps(range) }
+      let encodedTokens = tokens.lspEncoded
 
       req.reply(DocumentSemanticTokensResponse(data: encodedTokens))
     }