Translate the semantic token legend used by clangd to the semantic token legend used by SourceKit-LSP

clangd uses a completely different semantic token legend than SourceKit-LSP (it doesn’t even adhere to the ordering of the pre-defined token types) but we were passing index offsets from clangd through assuming that clangd uses the same legend, which was incorrect.

When retrieving semantic tokens from clangd, translate the semantic tokens from clangd’s legend to SourceKit-LSP’s legend.

rdar://129895062

Author: ahoppen

Documentation/LSP Extensions.md

@@ -35,8 +35,35 @@ Added field (this is an extension from clangd that SourceKit-LSP re-exposes):
 codeActions: CodeAction[]?
 ```
 
+## Semantic token modifiers
+
+Added the following cases from clangd
+
+```ts
+deduced = 'deduced'
+virtual = 'virtual'
+dependentName = 'dependentName'
+usedAsMutableReference = 'usedAsMutableReference'
+usedAsMutablePointer = 'usedAsMutablePointer'
+constructorOrDestructor = 'constructorOrDestructor'
+userDefined = 'userDefined'
+functionScope = 'functionScope'
+classScope = 'classScope'
+fileScope = 'fileScope'
+globalScope = 'globalScope'
+```
+
 ## Semantic token types
 
+Added the following cases from clangd
+
+```ts
+bracket = 'bracket'
+label = 'label'
+concept = 'concept'
+unknown = 'unknown'
+```
+
 Added case
 
 ```ts

Sources/LanguageServerProtocol/SupportTypes/SemanticTokenModifiers.swift

@@ -34,6 +34,19 @@ public struct SemanticTokenModifiers: OptionSet, Hashable, Sendable {
   public static let documentation = Self(rawValue: 1 << 8)
   public static let defaultLibrary = Self(rawValue: 1 << 9)
 
+  // The following are LSP extensions from clangd
+  public static let deduced = Self(rawValue: 1 << 10)
+  public static let virtual = Self(rawValue: 1 << 11)
+  public static let dependentName = Self(rawValue: 1 << 12)
+  public static let usedAsMutableReference = Self(rawValue: 1 << 13)
+  public static let usedAsMutablePointer = Self(rawValue: 1 << 14)
+  public static let constructorOrDestructor = Self(rawValue: 1 << 15)
+  public static let userDefined = Self(rawValue: 1 << 16)
+  public static let functionScope = Self(rawValue: 1 << 17)
+  public static let classScope = Self(rawValue: 1 << 18)
+  public static let fileScope = Self(rawValue: 1 << 19)
+  public static let globalScope = Self(rawValue: 1 << 20)
+
   public var name: String? {
     switch self {
     case .declaration: return "declaration"
@@ -46,13 +59,24 @@ public struct SemanticTokenModifiers: OptionSet, Hashable, Sendable {
     case .modification: return "modification"
     case .documentation: return "documentation"
     case .defaultLibrary: return "defaultLibrary"
+    case .deduced: return "deduced"
+    case .virtual: return "virtual"
+    case .dependentName: return "dependentName"
+    case .usedAsMutableReference: return "usedAsMutableReference"
+    case .usedAsMutablePointer: return "usedAsMutablePointer"
+    case .constructorOrDestructor: return "constructorOrDestructor"
+    case .userDefined: return "userDefined"
+    case .functionScope: return "functionScope"
+    case .classScope: return "classScope"
+    case .fileScope: return "fileScope"
+    case .globalScope: return "globalScope"
     default: return nil
     }
   }
 
   /// All available modifiers, in ascending order of the bit index
   /// they are represented with (starting at the rightmost bit).
-  public static let predefined: [Self] = [
+  public static let all: [Self] = [
     .declaration,
     .definition,
     .readonly,
@@ -63,5 +87,16 @@ public struct SemanticTokenModifiers: OptionSet, Hashable, Sendable {
     .modification,
     .documentation,
     .defaultLibrary,
+    .deduced,
+    .virtual,
+    .dependentName,
+    .usedAsMutableReference,
+    .usedAsMutablePointer,
+    .constructorOrDestructor,
+    .userDefined,
+    .functionScope,
+    .classScope,
+    .fileScope,
+    .globalScope,
   ]
 }

Sources/LanguageServerProtocol/SupportTypes/SemanticTokenTypes.swift

@@ -50,7 +50,18 @@ public struct SemanticTokenTypes: Hashable, Sendable {
   /// since 3.17.0
   public static let decorator = Self("decorator")
 
-  public static let predefined: [Self] = [
+  // The following are LSP extensions from clangd
+  public static let bracket = Self("bracket")
+  public static let label = Self("label")
+  public static let concept = Self("concept")
+  public static let unknown = Self("unknown")
+
+  /// An identifier that hasn't been further classified
+  ///
+  /// **(LSP Extension)**
+  public static let identifier = Self("identifier")
+
+  public static let all: [Self] = [
     .namespace,
     .type,
     .class,
@@ -73,5 +84,11 @@ public struct SemanticTokenTypes: Hashable, Sendable {
     .number,
     .regexp,
     .operator,
+    .decorator,
+    .bracket,
+    .label,
+    .concept,
+    .unknown,
+    .identifier,
   ]
 }

Sources/SourceKitLSP/CMakeLists.txt

@@ -11,6 +11,7 @@ add_library(SourceKitLSP STATIC
   MessageHandlingDependencyTracker.swift
   Rename.swift
   ResponseError+Init.swift
+  SemanticTokensLegend+SourceKitLSPLegend.swift
   SourceKitIndexDelegate.swift
   SourceKitLSPCommandMetadata.swift
   SourceKitLSPServer.swift
@@ -22,7 +23,9 @@ add_library(SourceKitLSP STATIC
   Workspace.swift
 )
 target_sources(SourceKitLSP PRIVATE
-  Clang/ClangLanguageService.swift)
+  Clang/ClangLanguageService.swift
+  Clang/SemanticTokenTranslator.swift
+)
 target_sources(SourceKitLSP PRIVATE
   Swift/AdjustPositionToStartOfIdentifier.swift
   Swift/CodeActions/AddDocumentation.swift

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -95,6 +95,9 @@ actor ClangLanguageService: LanguageService, MessageHandler {
   /// opened with.
   private var openDocuments: [DocumentURI: Language] = [:]
 
+  /// Type to map `clangd`'s semantic token legend to SourceKit-LSP's.
+  private var semanticTokensTranslator: SemanticTokensLegendTranslator? = nil
+
   /// While `clangd` is running, its PID.
   #if os(Windows)
   private var hClangd: HANDLE = INVALID_HANDLE_VALUE
@@ -412,6 +415,12 @@ extension ClangLanguageService {
 
     let result = try await clangd.send(initialize)
     self.capabilities = result.capabilities
+    if let legend = result.capabilities.semanticTokensProvider?.legend {
+      self.semanticTokensTranslator = SemanticTokensLegendTranslator(
+        clangdLegend: legend,
+        sourceKitLSPLegend: SemanticTokensLegend.sourceKitLSPLegend
+      )
+    }
     return result
   }
 
@@ -537,19 +546,50 @@ extension ClangLanguageService {
   }
 
   func documentSemanticTokens(_ req: DocumentSemanticTokensRequest) async throws -> DocumentSemanticTokensResponse? {
-    return try await forwardRequestToClangd(req)
+    guard var response = try await forwardRequestToClangd(req) else {
+      return nil
+    }
+    if let semanticTokensTranslator {
+      response.data = semanticTokensTranslator.translate(response.data)
+    }
+    return response
   }
 
   func documentSemanticTokensDelta(
     _ req: DocumentSemanticTokensDeltaRequest
   ) async throws -> DocumentSemanticTokensDeltaResponse? {
-    return try await forwardRequestToClangd(req)
+    guard var response = try await forwardRequestToClangd(req) else {
+      return nil
+    }
+    if let semanticTokensTranslator {
+      switch response {
+      case .tokens(var tokens):
+        tokens.data = semanticTokensTranslator.translate(tokens.data)
+        response = .tokens(tokens)
+      case .delta(var delta):
+        delta.edits = delta.edits.map {
+          var edit = $0
+          if let data = edit.data {
+            edit.data = semanticTokensTranslator.translate(data)
+          }
+          return edit
+        }
+        response = .delta(delta)
+      }
+    }
+    return response
   }
 
   func documentSemanticTokensRange(
     _ req: DocumentSemanticTokensRangeRequest
   ) async throws -> DocumentSemanticTokensResponse? {
-    return try await forwardRequestToClangd(req)
+    guard var response = try await forwardRequestToClangd(req) else {
+      return nil
+    }
+    if let semanticTokensTranslator {
+      response.data = semanticTokensTranslator.translate(response.data)
+    }
+    return response
   }
 
   func colorPresentation(_ req: ColorPresentationRequest) async throws -> [ColorPresentation] {

Sources/SourceKitLSP/Clang/SemanticTokenTranslator.swift

@@ -0,0 +1,136 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+
+/// `clangd` might use a different semantic token legend than SourceKit-LSP.
+///
+/// This type allows translation the semantic tokens from `clangd` into the token legend that is used by SourceKit-LSP.
+struct SemanticTokensLegendTranslator {
+  private enum Translation {
+    /// The token type or modifier from clangd does not exist in SourceKit-LSP
+    case doesNotExistInSourceKitLSP
+
+    /// The token type or modifier exists in SourceKit-LSP but it uses a different index. We need to translate the
+    /// clangd index to this SourceKit-LSP index.
+    case translation(UInt32)
+  }
+
+  /// For all token types whose representation in clang differs from the representation in SourceKit-LSP, maps the
+  /// index of that token type in clangd’s token type legend to the corresponding representation in SourceKit-LSP.
+  private let tokenTypeTranslations: [UInt32: Translation]
+
+  /// For all token modifiers whose representation in clang differs from the representation in SourceKit-LSP, maps the
+  /// index of that token modifier in clangd’s token type legend to the corresponding representation in SourceKit-LSP.
+  private let tokenModifierTranslations: [UInt32: Translation]
+
+  /// A bitmask that has all bits set to 1 that are used for clangd token modifiers which have a different
+  /// representation in SourceKit-LSP. If a token modifier does not have any bits set in common with this bitmask, no
+  /// token mapping needs to be performed.
+  private let tokenModifierTranslationBitmask: UInt32
+
+  /// For token types in clangd that do not exist in SourceKit-LSP's token legend, we need to map their token types to
+  /// some valid SourceKit-LSP token type. Use the token type with this index.
+  private let tokenTypeFallbackIndex: UInt32
+
+  init(clangdLegend: SemanticTokensLegend, sourceKitLSPLegend: SemanticTokensLegend) {
+    var tokenTypeTranslations: [UInt32: Translation] = [:]
+    for (index, tokenType) in clangdLegend.tokenTypes.enumerated() {
+      switch sourceKitLSPLegend.tokenTypes.firstIndex(of: tokenType) {
+      case index:
+        break
+      case nil:
+        logger.error("Token type '\(tokenType, privacy: .public)' from clangd does not exist in SourceKit-LSP's legend")
+        tokenTypeTranslations[UInt32(index)] = .doesNotExistInSourceKitLSP
+      case let sourceKitLSPIndex?:
+        logger.info(
+          "Token type '\(tokenType, privacy: .public)' from clangd at index \(index) translated to \(sourceKitLSPIndex)"
+        )
+        tokenTypeTranslations[UInt32(index)] = .translation(UInt32(sourceKitLSPIndex))
+      }
+    }
+    self.tokenTypeTranslations = tokenTypeTranslations
+
+    var tokenModifierTranslations: [UInt32: Translation] = [:]
+    for (index, tokenModifier) in clangdLegend.tokenModifiers.enumerated() {
+      switch sourceKitLSPLegend.tokenModifiers.firstIndex(of: tokenModifier) {
+      case index:
+        break
+      case nil:
+        logger.error(
+          "Token modifier '\(tokenModifier, privacy: .public)' from clangd does not exist in SourceKit-LSP's legend"
+        )
+        tokenModifierTranslations[UInt32(index)] = .doesNotExistInSourceKitLSP
+      case let sourceKitLSPIndex?:
+        logger.error(
+          "Token modifier '\(tokenModifier, privacy: .public)' from clangd at index \(index) translated to \(sourceKitLSPIndex)"
+        )
+        tokenModifierTranslations[UInt32(index)] = .translation(UInt32(sourceKitLSPIndex))
+      }
+    }
+    self.tokenModifierTranslations = tokenModifierTranslations
+
+    var tokenModifierTranslationBitmask: UInt32 = 0
+    for translatedIndex in tokenModifierTranslations.keys {
+      tokenModifierTranslationBitmask.setBitToOne(at: Int(translatedIndex))
+    }
+    self.tokenModifierTranslationBitmask = tokenModifierTranslationBitmask
+
+    self.tokenTypeFallbackIndex = UInt32(
+      sourceKitLSPLegend.tokenTypes.firstIndex(of: SemanticTokenTypes.unknown.name) ?? 0
+    )
+  }
+
+  func translate(_ data: [UInt32]) -> [UInt32] {
+    var data = data
+    // Translate token types, which are at offset n + 3.
+    for i in stride(from: 3, to: data.count, by: 5) {
+      switch tokenTypeTranslations[data[i]] {
+      case .doesNotExistInSourceKitLSP: data[i] = tokenTypeFallbackIndex
+      case .translation(let translatedIndex): data[i] = translatedIndex
+      case nil: break
+      }
+    }
+
+    // Translate token modifiers, which are at offset n + 4
+    for i in stride(from: 4, to: data.count, by: 5) {
+      guard data[i] & tokenModifierTranslationBitmask != 0 else {
+        // Fast path: There is nothing to translate
+        continue
+      }
+      var translatedModifiersBitmask: UInt32 = 0
+      for (clangdModifier, sourceKitLSPModifier) in tokenModifierTranslations {
+        guard data[i].hasBitSet(at: Int(clangdModifier)) else {
+          continue
+        }
+        switch sourceKitLSPModifier {
+        case .doesNotExistInSourceKitLSP: break
+        case .translation(let sourceKitLSPIndex): translatedModifiersBitmask.setBitToOne(at: Int(sourceKitLSPIndex))
+        }
+      }
+      data[i] = data[i] & ~tokenModifierTranslationBitmask | translatedModifiersBitmask
+    }
+
+    return data
+  }
+}
+
+fileprivate extension UInt32 {
+  mutating func hasBitSet(at index: Int) -> Bool {
+    return self & (1 << index) != 0
+  }
+
+  mutating func setBitToOne(at index: Int) {
+    self |= 1 << index
+  }
+}

Sources/SourceKitLSP/SemanticTokensLegend+SourceKitLSPLegend.swift

@@ -0,0 +1,31 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 LanguageServerProtocol
+
+extension SemanticTokenTypes {
+  // LSP doesn’t know about actors. Display actors as classes.
+  public static var actor: Self { Self.class }
+
+  /// Token types are looked up by index
+  public var tokenType: UInt32 {
+    UInt32(Self.all.firstIndex(of: self)!)
+  }
+}
+
+extension SemanticTokensLegend {
+  /// The semantic tokens legend that is used between SourceKit-LSP and the editor.
+  static let sourceKitLSPLegend = SemanticTokensLegend(
+    tokenTypes: SemanticTokenTypes.all.map(\.name),
+    tokenModifiers: SemanticTokenModifiers.all.compactMap(\.name)
+  )
+}