Merge pull request #1146 from gremlinflat/syntax-highlighting-tokens-refactor

Introduce ```SyntaxHighlightingTokens``` instead of ```[SyntaxHighlightingToken]```

Author: ahoppen

Sources/SKTestSupport/Array+SyntaxHighlightingToken.swift

@@ -106,11 +106,12 @@ public enum SkipUnless {
       let response = try unwrap(
         await testClient.send(DocumentSemanticTokensRequest(textDocument: TextDocumentIdentifier(uri)))
       )
-      let tokens = [SyntaxHighlightingToken](lspEncodedTokens: response.data)
+
+      let tokens = SyntaxHighlightingTokens(lspEncodedTokens: response.data)
 
       // If we don't have semantic token support in sourcekitd, the second token is an identifier based on the syntax
       // tree, not a property.
-      return tokens != [
+      return tokens.tokens != [
         SyntaxHighlightingToken(
           range: Position(line: 0, utf16index: 0)..<Position(line: 0, utf16index: 1),
           kind: .number,

Sources/SKTestSupport/SkipUnless.swift

@@ -39,6 +39,7 @@ target_sources(SourceKitLSP PRIVATE
   Swift/SwiftLanguageService.swift
   Swift/SymbolInfo.swift
   Swift/SyntaxHighlightingToken.swift
+  Swift/SyntaxHighlightingTokens.swift
   Swift/SyntaxHighlightingTokenParser.swift
   Swift/SyntaxTreeManager.swift
   Swift/VariableTypeInfo.swift

Sources/SourceKitLSP/CMakeLists.txt

@@ -19,7 +19,7 @@ import SwiftSyntax
 
 extension SwiftLanguageService {
   /// Requests the semantic highlighting tokens for the given snapshot from sourcekitd.
-  private func semanticHighlightingTokens(for snapshot: DocumentSnapshot) async throws -> [SyntaxHighlightingToken]? {
+  private func semanticHighlightingTokens(for snapshot: DocumentSnapshot) async throws -> SyntaxHighlightingTokens? {
     guard let buildSettings = await self.buildSettings(for: snapshot.uri), !buildSettings.isFallback else {
       return nil
     }
@@ -35,6 +35,7 @@ extension SwiftLanguageService {
     guard let skTokens: SKDResponseArray = dict[keys.semanticTokens] else {
       return nil
     }
+
     return SyntaxHighlightingTokenParser(sourcekitd: sourcekitd).parseTokens(skTokens, in: snapshot)
   }
 
@@ -49,7 +50,7 @@ extension SwiftLanguageService {
   private func mergedAndSortedTokens(
     for snapshot: DocumentSnapshot,
     in range: Range<Position>? = nil
-  ) async throws -> [SyntaxHighlightingToken] {
+  ) async throws -> SyntaxHighlightingTokens {
     async let tree = syntaxTreeManager.syntaxTree(for: snapshot)
     let semanticTokens = await orLog("Loading semantic tokens") { try await semanticHighlightingTokens(for: snapshot) }
 
@@ -59,11 +60,16 @@ extension SwiftLanguageService {
       } else {
         ByteSourceRange(offset: 0, length: await tree.totalLength.utf8Length)
       }
-    return
+
+    let tokens =
       await tree
       .classifications(in: range)
-      .flatMap({ $0.highlightingTokens(in: snapshot) })
-      .mergingTokens(with: semanticTokens ?? [])
+      .map { $0.highlightingTokens(in: snapshot) }
+      .reduce(into: SyntaxHighlightingTokens(tokens: [])) { $0.tokens += $1.tokens }
+
+    return
+      tokens
+      .mergingTokens(with: semanticTokens ?? SyntaxHighlightingTokens(tokens: []))
       .sorted { $0.start < $1.start }
   }
 
@@ -102,28 +108,30 @@ extension Range where Bound == Position {
 }
 
 extension SyntaxClassifiedRange {
-  fileprivate func highlightingTokens(in snapshot: DocumentSnapshot) -> [SyntaxHighlightingToken] {
+  fileprivate func highlightingTokens(in snapshot: DocumentSnapshot) -> SyntaxHighlightingTokens {
     guard let (kind, modifiers) = self.kind.highlightingKindAndModifiers else {
-      return []
+      return SyntaxHighlightingTokens(tokens: [])
     }
 
     guard
       let start: Position = snapshot.positionOf(utf8Offset: self.offset),
       let end: Position = snapshot.positionOf(utf8Offset: self.endOffset)
     else {
-      return []
+      return SyntaxHighlightingTokens(tokens: [])
     }
 
     let multiLineRange = start..<end
     let ranges = multiLineRange.splitToSingleLineRanges(in: snapshot)
 
-    return ranges.map {
+    let tokens = ranges.map {
       SyntaxHighlightingToken(
         range: $0,
         kind: kind,
         modifiers: modifiers
       )
     }
+
+    return SyntaxHighlightingTokens(tokens: tokens)
   }
 }
 

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -48,48 +48,6 @@ public struct SyntaxHighlightingToken: Hashable {
   }
 }
 
-extension Array where Element == SyntaxHighlightingToken {
-  /// The LSP representation of syntax highlighting tokens. Note that this
-  /// requires the tokens in this array to be sorted.
-  public var lspEncoded: [UInt32] {
-    var previous = Position(line: 0, utf16index: 0)
-    var rawTokens: [UInt32] = []
-    rawTokens.reserveCapacity(count * 5)
-
-    for token in self {
-      let lineDelta = token.start.line - previous.line
-      let charDelta =
-        token.start.utf16index - (
-          // The character delta is relative to the previous token's start
-          // 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),
-        UInt32(charDelta),
-        UInt32(token.utf16length),
-        token.kind.tokenType,
-        token.modifiers.rawValue,
-      ]
-    }
-
-    return rawTokens
-  }
-
-  /// Merges the tokens in this array into a new token array,
-  /// preferring the given array's tokens if duplicate ranges are
-  /// found.
-  public func mergingTokens(with other: [SyntaxHighlightingToken]) -> [SyntaxHighlightingToken] {
-    let otherRanges = Set(other.map(\.range))
-    return filter { !otherRanges.contains($0.range) } + other
-  }
-}
-
 extension SemanticTokenTypes {
   /// **(LSP Extension)**
   public static let identifier = Self("identifier")

Sources/SourceKitLSP/Swift/SyntaxHighlightingToken.swift

@@ -25,7 +25,7 @@ struct SyntaxHighlightingTokenParser {
   private func parseTokens(
     _ response: SKDResponseDictionary,
     in snapshot: DocumentSnapshot,
-    into tokens: inout [SyntaxHighlightingToken]
+    into tokens: inout SyntaxHighlightingTokens
   ) {
     let keys = sourcekitd.keys
 
@@ -52,7 +52,7 @@ struct SyntaxHighlightingTokenParser {
         let multiLineRange = start..<end
         let ranges = multiLineRange.splitToSingleLineRanges(in: snapshot)
 
-        tokens += ranges.map {
+        tokens.tokens += ranges.map {
           SyntaxHighlightingToken(
             range: $0,
             kind: kind,
@@ -70,16 +70,16 @@ struct SyntaxHighlightingTokenParser {
   private func parseTokens(
     _ response: SKDResponseArray,
     in snapshot: DocumentSnapshot,
-    into tokens: inout [SyntaxHighlightingToken]
+    into tokens: inout SyntaxHighlightingTokens
   ) {
     response.forEach { (_, value) in
       parseTokens(value, in: snapshot, into: &tokens)
       return true
     }
   }
 
-  func parseTokens(_ response: SKDResponseArray, in snapshot: DocumentSnapshot) -> [SyntaxHighlightingToken] {
-    var tokens: [SyntaxHighlightingToken] = []
+  func parseTokens(_ response: SKDResponseArray, in snapshot: DocumentSnapshot) -> SyntaxHighlightingTokens {
+    var tokens: SyntaxHighlightingTokens = SyntaxHighlightingTokens(tokens: [])
     parseTokens(response, in: snapshot, into: &tokens)
     return tokens
   }

Sources/SourceKitLSP/Swift/SyntaxHighlightingTokenParser.swift

@@ -0,0 +1,115 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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 LSPLogging
+import LanguageServerProtocol
+import SourceKitD
+
+/// A wrapper around an array of syntax highlighting tokens.
+public struct SyntaxHighlightingTokens {
+  public var tokens: [SyntaxHighlightingToken]
+
+  public init(tokens: [SyntaxHighlightingToken]) {
+    self.tokens = tokens
+  }
+
+  /// The LSP representation of syntax highlighting tokens. Note that this
+  /// requires the tokens in this array to be sorted.
+  public var lspEncoded: [UInt32] {
+    var previous = Position(line: 0, utf16index: 0)
+    var rawTokens: [UInt32] = []
+    rawTokens.reserveCapacity(tokens.count * 5)
+
+    for token in self.tokens {
+      let lineDelta = token.start.line - previous.line
+      let charDelta =
+        token.start.utf16index - (
+          // The character delta is relative to the previous token's start
+          // 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),
+        UInt32(charDelta),
+        UInt32(token.utf16length),
+        token.kind.tokenType,
+        token.modifiers.rawValue,
+      ]
+    }
+
+    return rawTokens
+  }
+
+  /// Merges the tokens in this array into a new token array,
+  /// preferring the given array's tokens if duplicate ranges are
+  /// found.
+  public func mergingTokens(with other: SyntaxHighlightingTokens) -> SyntaxHighlightingTokens {
+    let otherRanges = Set(other.tokens.map(\.range))
+    return SyntaxHighlightingTokens(tokens: tokens.filter { !otherRanges.contains($0.range) } + other.tokens)
+  }
+
+  public func mergingTokens(with other: [SyntaxHighlightingToken]) -> SyntaxHighlightingTokens {
+    let otherRanges = Set(other.map(\.range))
+    return SyntaxHighlightingTokens(tokens: tokens.filter { !otherRanges.contains($0.range) } + other)
+  }
+
+  /// Sorts the tokens in this array by their start position.
+  public func sorted(_ areInIncreasingOrder: (SyntaxHighlightingToken, SyntaxHighlightingToken) -> Bool)
+    -> SyntaxHighlightingTokens
+  {
+    SyntaxHighlightingTokens(tokens: tokens.sorted(by: areInIncreasingOrder))
+  }
+}
+
+extension SyntaxHighlightingTokens {
+  /// Decodes the LSP representation of syntax highlighting tokens
+  public init(lspEncodedTokens rawTokens: [UInt32]) {
+    self.init(tokens: [])
+    assert(rawTokens.count.isMultiple(of: 5))
+    self.tokens.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]
+
+      current.line += lineDelta
+
+      if lineDelta == 0 {
+        current.utf16index += charDelta
+      } else {
+        current.utf16index = charDelta
+      }
+
+      let kind = SemanticTokenTypes.all[Int(rawKind)]
+      let modifiers = SemanticTokenModifiers(rawValue: rawModifiers)
+
+      self.tokens.append(
+        SyntaxHighlightingToken(
+          start: current,
+          utf16length: length,
+          kind: kind,
+          modifiers: modifiers
+        )
+      )
+    }
+  }
+}