Filter enclosing syntactic tokens during semantic token merging

Sources/SwiftLanguageService/SemanticTokens.swift

@@ -62,15 +62,15 @@ extension SwiftLanguageService {
 
     async let semanticTokens =
       await orLog("Loading semantic tokens") { try await semanticHighlightingTokens(for: snapshot) }
-      ?? SyntaxHighlightingTokens(tokens: [])
+      ?? SyntaxHighlightingTokens(sortedTokens: [])
 
     let tokens =
       if self.options.reportSyntacticHighlightInSemanticTokensOrDefault {
         try await syntacticTokens(for: snapshot, in: range).mergingTokens(with: semanticTokens)
       } else {
         await semanticTokens
       }
-    return tokens.sorted({ $0.start < $1.start })
+    return tokens
   }
 
   private func syntacticTokens(
@@ -88,11 +88,11 @@ extension SwiftLanguageService {
         tree.range
       }
 
-    return
-      tree
-      .classifications(in: range)
-      .map { $0.highlightingTokens(in: snapshot) }
-      .reduce(into: SyntaxHighlightingTokens(tokens: [])) { $0.tokens += $1.tokens }
+    return SyntaxHighlightingTokens(
+      sortedTokens: tree
+        .classifications(in: range)
+        .flatMap { $0.highlightingTokens(in: snapshot).tokens }
+    )
   }
 
   package func documentSemanticTokens(
@@ -126,7 +126,7 @@ extension SwiftLanguageService {
 extension SyntaxClassifiedRange {
   fileprivate func highlightingTokens(in snapshot: DocumentSnapshot) -> SyntaxHighlightingTokens {
     guard let (kind, modifiers) = self.kind.highlightingKindAndModifiers else {
-      return SyntaxHighlightingTokens(tokens: [])
+      return SyntaxHighlightingTokens(sortedTokens: [])
     }
 
     let multiLineRange = snapshot.positionRange(of: self.range)
@@ -140,7 +140,7 @@ extension SyntaxClassifiedRange {
       )
     }
 
-    return SyntaxHighlightingTokens(tokens: tokens)
+    return SyntaxHighlightingTokens(sortedTokens: tokens)
   }
 }
 

Sources/SwiftLanguageService/SyntaxHighlightingTokenParser.swift

@@ -27,7 +27,7 @@ struct SyntaxHighlightingTokenParser {
   private func parseTokens(
     _ response: SKDResponseDictionary,
     in snapshot: DocumentSnapshot,
-    into tokens: inout SyntaxHighlightingTokens
+    into tokens: inout [SyntaxHighlightingToken]
   ) {
     let keys = sourcekitd.keys
 
@@ -52,7 +52,7 @@ struct SyntaxHighlightingTokenParser {
       let multiLineRange = snapshot.positionOf(utf8Offset: offset)..<snapshot.positionOf(utf8Offset: offset + length)
       let ranges = multiLineRange.splitToSingleLineRanges(in: snapshot)
 
-      tokens.tokens += ranges.map {
+      tokens += ranges.map {
         SyntaxHighlightingToken(
           range: $0,
           kind: kind,
@@ -69,7 +69,7 @@ struct SyntaxHighlightingTokenParser {
   private func parseTokens(
     _ response: SKDResponseArray,
     in snapshot: DocumentSnapshot,
-    into tokens: inout SyntaxHighlightingTokens
+    into tokens: inout [SyntaxHighlightingToken]
   ) {
     // swift-format-ignore: ReplaceForEachWithForLoop
     // Reference is to `SKDResponseArray.forEach`, not `Array.forEach`.
@@ -80,9 +80,9 @@ struct SyntaxHighlightingTokenParser {
   }
 
   func parseTokens(_ response: SKDResponseArray, in snapshot: DocumentSnapshot) -> SyntaxHighlightingTokens {
-    var tokens: SyntaxHighlightingTokens = SyntaxHighlightingTokens(tokens: [])
+    var tokens: [SyntaxHighlightingToken] = []
     parseTokens(response, in: snapshot, into: &tokens)
-    return tokens
+    return SyntaxHighlightingTokens(tokens: tokens)
   }
 
   private func parseKindAndModifiers(

Sources/SwiftLanguageService/SyntaxHighlightingTokens.swift

@@ -17,10 +17,26 @@ import SourceKitLSP
 
 /// A wrapper around an array of syntax highlighting tokens.
 package struct SyntaxHighlightingTokens: Sendable {
-  package var tokens: [SyntaxHighlightingToken]
 
+  /// Syntax highlighting tokens sorted by their start position.
+  package let tokens: [SyntaxHighlightingToken]
+
+  /// Creates a syntax highlighting token collection from an sorted array.
+  ///
+  /// - Parameter sortedTokens: Syntax highlighting tokens sorted by their start position.
+  package init(sortedTokens: [SyntaxHighlightingToken]) {
+    assert(
+      zip(sortedTokens, sortedTokens.dropFirst()).allSatisfy { $0.start <= $1.start },
+      "Tokens must always be sorted by their start position"
+    )
+    self.tokens = sortedTokens
+  }
+
+  /// Creates a syntax highlighting token collection from a potentially unsorted array.
+  ///
+  /// - Parameter tokens: Syntax highlighting tokens that will be sorted by their start position.
   package init(tokens: [SyntaxHighlightingToken]) {
-    self.tokens = tokens
+    self.init(sortedTokens: tokens.sorted { $0.start < $1.start })
   }
 
   /// The LSP representation of syntax highlighting tokens. Note that this
@@ -56,32 +72,50 @@ package struct SyntaxHighlightingTokens: Sendable {
   }
 
   /// Merges the tokens in this array into a new token array,
-  /// preferring the given array's tokens if duplicate ranges are
+  /// preferring the given array's tokens if overlapping ranges are
   /// found.
   package func mergingTokens(with other: SyntaxHighlightingTokens) -> SyntaxHighlightingTokens {
-    let otherRanges = Set(other.tokens.map(\.range))
-    return SyntaxHighlightingTokens(tokens: tokens.filter { !otherRanges.contains($0.range) } + other.tokens)
-  }
+    var merged: [SyntaxHighlightingToken] = []
+    merged.reserveCapacity(tokens.count + other.tokens.count)
 
-  package func mergingTokens(with other: [SyntaxHighlightingToken]) -> SyntaxHighlightingTokens {
-    let otherRanges = Set(other.map(\.range))
-    return SyntaxHighlightingTokens(tokens: tokens.filter { !otherRanges.contains($0.range) } + other)
-  }
+    var selfIterator = tokens.makeIterator()
+    var otherIterator = other.tokens.makeIterator()
+
+    var currentToken = selfIterator.next()
+    var currentOtherToken = otherIterator.next()
+
+    while let token = currentToken, let otherToken = currentOtherToken {
+      if token.range.overlaps(otherToken.range) {
+        currentToken = selfIterator.next()
+      } else if token.start < otherToken.start {
+        merged.append(token)
+        currentToken = selfIterator.next()
+      } else {
+        merged.append(otherToken)
+        currentOtherToken = otherIterator.next()
+      }
+    }
+
+    while let token = currentToken {
+      merged.append(token)
+      currentToken = selfIterator.next()
+    }
+
+    while let otherToken = currentOtherToken {
+      merged.append(otherToken)
+      currentOtherToken = otherIterator.next()
+    }
 
-  /// Sorts the tokens in this array by their start position.
-  package func sorted(
-    _ areInIncreasingOrder: (SyntaxHighlightingToken, SyntaxHighlightingToken) -> Bool
-  ) -> SyntaxHighlightingTokens {
-    SyntaxHighlightingTokens(tokens: tokens.sorted(by: areInIncreasingOrder))
+    return SyntaxHighlightingTokens(sortedTokens: merged)
   }
 }
 
 extension SyntaxHighlightingTokens {
   /// Decodes the LSP representation of syntax highlighting tokens
   package init(lspEncodedTokens rawTokens: [UInt32]) {
-    self.init(tokens: [])
     assert(rawTokens.count.isMultiple(of: 5))
-    self.tokens.reserveCapacity(rawTokens.count / 5)
+    var parsedTokens: [SyntaxHighlightingToken] = []
+    parsedTokens.reserveCapacity(rawTokens.count / 5)
 
     var current = Position(line: 0, utf16index: 0)
 
@@ -103,7 +137,7 @@ extension SyntaxHighlightingTokens {
       let kind = SemanticTokenTypes.all[Int(rawKind)]
       let modifiers = SemanticTokenModifiers(rawValue: rawModifiers)
 
-      self.tokens.append(
+      parsedTokens.append(
         SyntaxHighlightingToken(
           start: current,
           utf16length: length,
@@ -112,5 +146,6 @@ extension SyntaxHighlightingTokens {
         )
       )
     }
+    self.init(sortedTokens: parsedTokens)
   }
 }

Tests/SourceKitLSPTests/SemanticTokensTests.swift

@@ -24,7 +24,7 @@ private typealias Token = SyntaxHighlightingToken
 
 final class SemanticTokensTests: SourceKitLSPTestCase {
   func testIntArrayCoding() async throws {
-    let tokens = SyntaxHighlightingTokens(tokens: [
+    let tokens = SyntaxHighlightingTokens(sortedTokens: [
       Token(
         start: Position(line: 2, utf16index: 3),
         utf16length: 5,
@@ -1048,6 +1048,25 @@ final class SemanticTokensTests: SourceKitLSPTestCase {
     )
     try await fulfillmentOfOrThrow(receivedSemanticTokensResponse)
   }
+
+  func testAttachedMacroSemanticTokens() async throws {
+    try await assertSemanticTokens(
+      markedContents: """
+        @attached(accessor) @attached(peer, names: prefixed(`$`))
+        public macro TaskLocal() = #externalMacro(module: "SwiftMacros", type: "TaskLocal")
+
+        class C {
+            1️⃣@2️⃣TaskLocal3️⃣
+            static var taskLocalLogger: 4️⃣String = ""
+        }
+        """,
+      range: ("1️⃣", "3️⃣"),
+      expected: [
+        TokenSpec(marker: "2️⃣", length: 9, kind: .macro, isSourceKit: true),
+        TokenSpec(marker: "4️⃣", length: 6, kind: .struct, modifiers: .defaultLibrary, isSourceKit: true),
+      ]
+    )
+  }
 }
 
 private struct TokenSpec {