Address PR suggestions regarding semantic tokens

- Remove token bounding logic in DocumentManager
  Also address some stylistic suggestions.
- Make hasSemanticTokensRegistration take only a single language
- Make lspName internal and add _lspName wrappers for testing
- Assert that tokens are sorted while encoding
- Compute function token length in terms of UTF-16 code units
  ...these weren't handled correctly before for emoji.
- Let splitToSingleLineTokens return a non-optional

Author: fwcd

Sources/SourceKitLSP/CapabilityRegistry.swift

@@ -143,8 +143,8 @@ public final class CapabilityRegistry {
   }
 
   /// Checks whether a registration for semantic tokens for the given languages exists.
-  public func hasSemanticTokensRegistration(for languages: [Language]) -> Bool {
-    registration(for: languages, in: semanticTokens) != nil
+  public func hasSemanticTokensRegistration(for language: Language) -> Bool {
+    registration(for: [language], in: semanticTokens) != nil
   }
 
   private func documentSelector(for langauges: [Language]) -> DocumentSelector {

Sources/SourceKitLSP/DocumentManager.swift

@@ -23,7 +23,7 @@ public struct DocumentTokens {
   /// Semantic tokens, e.g. variable references, type references, ...
   public var semantic: [SyntaxHighlightingToken] = []
 
-  public var merged: [SyntaxHighlightingToken] {
+  private var merged: [SyntaxHighlightingToken] {
     [lexical, syntactic, semantic].reduce([]) { $0.mergingTokens(with: $1) }
   }
   public var mergedAndSorted: [SyntaxHighlightingToken] {
@@ -174,19 +174,11 @@ public final class DocumentManager {
           let lastLineLengthDelta = newLines.last!.count - lastLineReplaceLength
           let lineDelta = newLines.count - previousLineCount
 
-          func isTokenBounding(character: Character) -> Bool {
-            character.isWhitespace || character.isPunctuation || character.isSymbol
-          }
-
-          func update(tokens: inout [SyntaxHighlightingToken]) {
+          document.latestTokens.withMutableTokensOfEachKind { tokens in
             tokens = Array(tokens.lazy
               .filter {
-                // Only keep tokens that don't overlap with or are directly
-                // adjacent to the edit range and also are adjacent to a
-                // token-bounding character.
-                $0.start > range.upperBound || range.lowerBound > $0.sameLineEnd
-                || ($0.start == range.upperBound && (edit.text.first.map(isTokenBounding(character:)) ?? true))
-                || ($0.sameLineEnd == range.lowerBound && (edit.text.last.map(isTokenBounding(character:)) ?? true))
+                // Only keep tokens that don't overlap with the edit range
+                $0.start >= range.upperBound || range.lowerBound >= $0.sameLineEnd
               }
               .map {
                 // Shift tokens after the edit range
@@ -201,8 +193,6 @@ public final class DocumentManager {
                 return token
               })
           }
-
-          document.latestTokens.withMutableTokensOfEachKind(update(tokens:))
         } else {
           // Full text replacement.
           document.latestLineTable = LineTable(edit.text)
@@ -259,7 +249,7 @@ public final class DocumentManager {
   ///
   /// - parameter uri: The URI of the document to be updated
   /// - parameter range: The range to replace tokens in (nil means the entire document)
-  /// - parameter tokens: The tokens to be added
+  /// - parameter newTokens: The tokens to be added
   @discardableResult
   public func replaceLexicalTokens(
     _ uri: DocumentURI,
@@ -271,17 +261,15 @@ public final class DocumentManager {
         throw Error.missingDocument(uri)
       }
 
-      if !newTokens.isEmpty {
-        if let range = range {
-          document.latestTokens.lexical.removeAll {
-            // Remove overlapping or bounding tokens
-            $0.start <= range.upperBound && range.lowerBound <= $0.sameLineEnd
-          }
-        }
-
-        document.latestTokens.lexical += newTokens
+      // Remove all tokens in `range` (or the entire document if `range` is `nil`)
+      document.latestTokens.lexical.removeAll { token in
+        range.map {
+          token.start <= $0.upperBound && $0.lowerBound <= token.sameLineEnd
+        } ?? true
       }
 
+      document.latestTokens.lexical += newTokens
+
       return document.latestSnapshot
     }
   }

Sources/SourceKitLSP/SourceKitServer.swift

@@ -861,7 +861,7 @@ extension SourceKitServer {
       return false
     }
     let language = snapshot.document.language
-    return workspace.capabilityRegistry.hasSemanticTokensRegistration(for: [language])
+    return workspace.capabilityRegistry.hasSemanticTokensRegistration(for: language)
   }
 
   func documentSemanticTokens(

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -221,7 +221,6 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
 
     do {
       try documentManager.replaceSemanticTokens(uri, tokens: tokens)
-      
     } catch {
       log("updating semantic tokens for \(uri) failed: \(error)", level: .warning)
     }

Sources/SourceKitLSP/Swift/SyntaxHighlightingToken.swift

@@ -32,7 +32,6 @@ public struct SyntaxHighlightingToken: Hashable {
   }
 
   public init(
-    name: String? = nil,
     start: Position,
     length: Int,
     kind: Kind,
@@ -45,9 +44,9 @@ public struct SyntaxHighlightingToken: Hashable {
   }
 
   /// Splits a potentially multi-line token to multiple single-line tokens.
-  public func splitToSingleLineTokens(in snapshot: DocumentSnapshot) -> [Self]? {
+  public func splitToSingleLineTokens(in snapshot: DocumentSnapshot) -> [Self] {
     guard let startIndex = snapshot.index(of: start) else {
-      return nil
+      fatalError("Token \(self) begins outside of the document")
     }
 
     let endIndex = snapshot.text.index(startIndex, offsetBy: length)
@@ -104,8 +103,7 @@ public struct SyntaxHighlightingToken: Hashable {
     case `operator`
 
     /// The name of the token type used by LSP.
-    /// **Public for testing.**
-    public var lspName: String {
+    var lspName: String {
       switch self {
       case .namespace: return "namespace"
       case .type: return "type"
@@ -131,6 +129,11 @@ public struct SyntaxHighlightingToken: Hashable {
       case .operator: return "operator"
       }
     }
+
+    /// **Public for testing**
+    public var _lspName: String {
+      lspName
+    }
   }
 
   /// Additional metadata about a token. Similar to `Kind`, the raw
@@ -169,8 +172,7 @@ public struct SyntaxHighlightingToken: Hashable {
     /// The name of the modifier used by LSP, if this
     /// is a single modifier. Note that every modifier
     /// in `allCases` must have an associated `lspName`.
-    /// **Public for testing.**
-    public var lspName: String? {
+    var lspName: String? {
       switch self {
       case .declaration: return "declaration"
       case .definition: return "definition"
@@ -186,6 +188,11 @@ public struct SyntaxHighlightingToken: Hashable {
       }
     }
 
+    /// **Public for testing**
+    public var _lspName: String? {
+      lspName
+    }
+
     public init(rawValue: UInt32) {
       self.rawValue = rawValue
     }
@@ -207,6 +214,11 @@ extension Array where Element == SyntaxHighlightingToken {
         // only if the token is on the previous token's line.
         previous.line == token.start.line ? previous.utf16index : 0
       )
+
+      // We assert that the tokens are actually sorted
+      assert(lineDelta >= 0)
+      assert(charDelta >= 0)
+
       previous = token.start
       rawTokens += [
         UInt32(lineDelta),
@@ -259,7 +271,7 @@ struct SyntaxHighlightingTokenParser {
       if useName && [.function, .method, .enumMember].contains(kind) && modifiers.contains(.declaration),
          let name: String = response[keys.name],
          name.contains("("),
-         let funcNameLength: Int = name.split(separator: "(").first?.count {
+         let funcNameLength: Int = name.split(separator: "(").first?.utf16.count {
         length = funcNameLength
       }
 
@@ -277,9 +289,7 @@ struct SyntaxHighlightingTokenParser {
         modifiers: modifiers
       )
 
-      if let newTokens = multiLineToken.splitToSingleLineTokens(in: snapshot) {
-        tokens += newTokens
-      }
+      tokens += multiLineToken.splitToSingleLineTokens(in: snapshot)
     }
 
     if let substructure: SKDResponseArray = response[keys.substructure] {

Tests/SourceKitLSPTests/SemanticTokensTests.swift

@@ -51,8 +51,8 @@ final class SemanticTokensTests: XCTestCase {
               range: .bool(true),
               full: .bool(true)
             ),
-            tokenTypes: Token.Kind.allCases.map(\.lspName),
-            tokenModifiers: Token.Modifiers.allCases.map { $0.lspName! },
+            tokenTypes: Token.Kind.allCases.map(\._lspName),
+            tokenModifiers: Token.Modifiers.allCases.map { $0._lspName! },
             formats: [.relative]
           )
         )
@@ -320,6 +320,15 @@ final class SemanticTokensTests: XCTestCase {
     ])
   }
 
+  func testSemanticTokensForFunctionSignaturesWithEmoji() {
+    let text = "func 👍abc() {}"
+    let tokens = performSemanticTokensRequest(text: text)
+    XCTAssertEqual(tokens, [
+      Token(start: Position(line: 0, utf16index: 0), length: 4, kind: .keyword),
+      Token(start: Position(line: 0, utf16index: 5), length: 5, kind: .function, modifiers: .declaration),
+    ])
+  }
+
   func testSemanticTokensForStaticMethods() {
     let text = """
     class X {

Tests/SourceKitLSPTests/XCTestManifests.swift

@@ -165,6 +165,7 @@ extension SemanticTokensTests {
       ("testSemanticTokens", testSemanticTokens),
       ("testSemanticTokensForEnumMembers", testSemanticTokensForEnumMembers),
       ("testSemanticTokensForFunctionSignatures", testSemanticTokensForFunctionSignatures),
+      ("testSemanticTokensForFunctionSignaturesWithEmoji", testSemanticTokensForFunctionSignaturesWithEmoji),
       ("testSemanticTokensForProtocols", testSemanticTokensForProtocols),
       ("testSemanticTokensForStaticMethods", testSemanticTokensForStaticMethods),
       ("testSyntacticTokens", testSyntacticTokens),