Keep track of semantic tokens in a SemanticTokensManager

Storing the semantic tokens inside `Document` was an anti-pattern because the semantic tokens only applied to Swift and were also being updated while the document contents themselves stayed constant.

Instead, we should store the semantic tokens in a separate `SemanticTokensManager` that only exists in the `SwiftLanguageServer` and has the sole responsibility of tracking semantic tokens.

Author: ahoppen

Sources/SourceKitLSP/CMakeLists.txt

@@ -2,7 +2,6 @@
 add_library(SourceKitLSP STATIC
   CapabilityRegistry.swift
   DocumentManager.swift
-  DocumentTokens.swift
   IndexStoreDB+MainFilesProvider.swift
   RangeAdjuster.swift
   Sequence+AsyncMap.swift
@@ -26,6 +25,7 @@ target_sources(SourceKitLSP PRIVATE
   Swift/SemanticRefactorCommand.swift
   Swift/SemanticRefactoring.swift
   Swift/SemanticTokens.swift
+  Swift/SemanticTokensManager.swift
   Swift/SourceKitD+ResponseError.swift
   Swift/SwiftCommand.swift
   Swift/SwiftLanguageServer.swift

Sources/SourceKitLSP/DocumentManager.swift

@@ -41,10 +41,6 @@ public struct DocumentSnapshot: Identifiable {
   public let id: ID
   public let language: Language
   public let 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 let tokens: DocumentTokens
 
   public var uri: DocumentURI { id.uri }
   public var version: Int { id.version }
@@ -54,13 +50,11 @@ public struct DocumentSnapshot: Identifiable {
     uri: DocumentURI,
     language: Language,
     version: Int,
-    lineTable: LineTable,
-    tokens: DocumentTokens
+    lineTable: LineTable
   ) {
     self.id = ID(uri: uri, version: version)
     self.language = language
     self.lineTable = lineTable
-    self.tokens = tokens
   }
 
   func index(of pos: Position) -> String.Index? {
@@ -73,14 +67,12 @@ public final class Document {
   public let language: Language
   var latestVersion: Int
   var latestLineTable: LineTable
-  var latestTokens: DocumentTokens
 
   init(uri: DocumentURI, language: Language, version: Int, text: String) {
     self.uri = uri
     self.language = language
     self.latestVersion = version
     self.latestLineTable = LineTable(text)
-    self.latestTokens = DocumentTokens()
   }
 
   /// **Not thread safe!** Use `DocumentManager.latestSnapshot` instead.
@@ -89,8 +81,7 @@ public final class Document {
       uri: self.uri,
       language: self.language,
       version: latestVersion,
-      lineTable: latestLineTable,
-      tokens: latestTokens
+      lineTable: latestLineTable
     )
   }
 }
@@ -178,23 +169,9 @@ public final class DocumentManager {
             toLine: range.upperBound.line,
             utf16Offset: range.upperBound.utf16index,
             with: edit.text)
-          
-          // Remove all tokens in the updated range and shift later ones.
-          let rangeAdjuster = RangeAdjuster(edit: edit)!
-
-          document.latestTokens.semantic = document.latestTokens.semantic.compactMap {
-            var token = $0
-            if let adjustedRange = rangeAdjuster.adjust(token.range) {
-              token.range = adjustedRange
-              return token
-            } else {
-              return nil
-            }
-          }
         } else {
           // Full text replacement.
           document.latestLineTable = LineTable(edit.text)
-          document.latestTokens = DocumentTokens()
         }
       }
 
@@ -206,23 +183,6 @@ public final class DocumentManager {
     }
   }
 
-  /// Updates the tokens in a document.
-  ///
-  /// - parameter uri: The URI of the document to be updated
-  /// - parameter tokens: The new tokens for the document
-  @discardableResult
-  public func updateTokens(_ uri: DocumentURI, tokens: DocumentTokens) throws -> DocumentSnapshot {
-    return try queue.sync {
-      guard let document = documents[uri] else {
-        throw Error.missingDocument(uri)
-      }
-
-      document.latestTokens = tokens
-
-      return document.latestSnapshot
-    }
-  }
-
   public func latestSnapshot(_ uri: DocumentURI) -> DocumentSnapshot? {
     return queue.sync {
       guard let document = documents[uri] else {

Sources/SourceKitLSP/DocumentTokens.swift

@@ -30,12 +30,13 @@ extension SwiftLanguageServer {
     in range: Range<Position>? = nil
   ) async -> [SyntaxHighlightingToken] {
     let tree = await syntaxTreeManager.syntaxTree(for: snapshot)
+    let semanticTokens = await semanticTokensManager.semanticTokens(for: snapshot.id)
     let range = range.flatMap({ $0.byteSourceRange(in: snapshot) })
              ?? ByteSourceRange(offset: 0, length: tree.totalLength.utf8Length)
     return tree
       .classifications(in: range)
       .flatMap({ $0.highlightingTokens(in: snapshot) })
-      .mergingTokens(with: snapshot.tokens.semantic)
+      .mergingTokens(with: semanticTokens ?? [])
       .sorted { $0.start < $1.start }
   }
 

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -0,0 +1,90 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+
+/// Keeps track of the semantic tokens that sourcekitd has sent us for given
+/// document snapshots.
+actor SemanticTokensManager {
+  private var semanticTokens: [DocumentSnapshot.ID: [SyntaxHighlightingToken]] = [:]
+
+  /// The semantic tokens for the given snapshot or `nil` if no semantic tokens
+  /// have been computed yet.
+  func semanticTokens(for snapshotID: DocumentSnapshot.ID) -> [SyntaxHighlightingToken]? {
+    return semanticTokens[snapshotID]
+  }
+
+  /// Set the semantic tokens that sourcekitd has sent us for the given document
+  /// snapshot.
+  ///
+  /// This discards any semantic tokens for any older versions of this document.
+  func setSemanticTokens(for snapshotID: DocumentSnapshot.ID, semanticTokens tokens: [SyntaxHighlightingToken]) {
+    semanticTokens[snapshotID] = tokens
+    // Delete semantic tokens for older versions of this document.
+    for key in semanticTokens.keys {
+      if key < snapshotID {
+        semanticTokens[key] = nil
+      }
+    }
+  }
+
+  /// If we have semantic tokens for `preEditSnapshotID`, shift the tokens
+  /// according to `edits` and store these shifted results for `postEditSnapshot`.
+  ///
+  /// This allows us to maintain semantic tokens after an edit for all the
+  /// non-edited regions.
+  ///
+  /// - Note: The semantic tokens stored from this edit might not be correct if
+  ///   the edits affect semantic highlighting for tokens out of the edit region.
+  ///   These will be updated when sourcekitd sends us new semantic tokens,
+  ///   which are stored in `SemanticTokensManager` by calling `setSemanticTokens`.
+  func registerEdit(
+    preEditSnapshot preEditSnapshotID: DocumentSnapshot.ID,
+    postEditSnapshot postEditSnapshotID: DocumentSnapshot.ID,
+    edits: [TextDocumentContentChangeEvent]
+  ) {
+    guard var semanticTokens = semanticTokens(for: preEditSnapshotID) else {
+      return
+    }
+    for edit in edits {
+      // Remove all tokens in the updated range and shift later ones.
+      guard let rangeAdjuster = RangeAdjuster(edit: edit) else {
+        // We have a full document edit and can't update semantic tokens
+        return
+      }
+
+      semanticTokens = semanticTokens.compactMap {
+        var token = $0
+        if let adjustedRange = rangeAdjuster.adjust(token.range) {
+          token.range = adjustedRange
+          return token
+        } else {
+          return nil
+        }
+      }
+    }
+    setSemanticTokens(for: postEditSnapshotID, semanticTokens: semanticTokens)
+  }
+
+  /// Discard any semantic tokens for documents with the given URI.
+  ///
+  /// This should be called when a document is being closed and the semantic
+  /// tokens are thus no longer needed.
+  func discaredSemanticTokens(for document: DocumentURI) {
+    // Delete semantic tokens for older versions of this document.
+    for key in semanticTokens.keys {
+      if key.uri == document {
+        semanticTokens[key] = nil
+      }
+    }
+  }
+}

Sources/SourceKitLSP/Swift/SemanticTokensManager.swift

@@ -126,6 +126,7 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
   var currentCompletionSession: CodeCompletionSession? = nil
 
   let syntaxTreeManager = SyntaxTreeManager()
+  let semanticTokensManager = SemanticTokensManager()
 
   nonisolated var keys: sourcekitd_keys { return sourcekitd.keys }
   nonisolated var requests: sourcekitd_requests { return sourcekitd.requests }
@@ -198,37 +199,21 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
     self.stateChangeHandlers.append(handler)
   }
 
-  /// Updates the semantic tokens for the given `snapshot`.
-  private func updateSemanticTokens(
-    response: SKDResponseDictionary,
-    for snapshot: DocumentSnapshot
-  ) {
-    let docTokens = updatedSemanticTokens(response: response, for: snapshot)
-
-    do {
-      try documentManager.updateTokens(snapshot.uri, tokens: docTokens)
-    } catch {
-      log("Updating semantic tokens failed: \(error)", level: .warning)
-    }
-  }
-
-  /// Returns the updated semantic tokens for the given `snapshot`.
-  private func updatedSemanticTokens(
-    response: SKDResponseDictionary,
+  /// Returns the semantic tokens in `response` for the given `snapshot`.
+  private func semanticTokens(
+    of response: SKDResponseDictionary,
     for snapshot: DocumentSnapshot
-  ) -> DocumentTokens {
+  ) -> [SyntaxHighlightingToken]? {
     logExecutionTime(level: .debug) {
-      var docTokens = snapshot.tokens
-
       if let skTokens: SKDResponseArray = response[keys.annotations] {
         let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd)
         var tokens: [SyntaxHighlightingToken] = []
         tokenParser.parseTokens(skTokens, in: snapshot, into: &tokens)
 
-        docTokens.semantic = tokens
+        return tokens
       }
 
-      return docTokens
+      return nil
     }
   }
 
@@ -344,9 +329,9 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
 
     if let dict = try? self.sourcekitd.sendSync(req) {
       let isSemaStage = dict[keys.diagnostic_stage] as sourcekitd_uid_t? == sourcekitd.values.diag_stage_sema
-      if isSemaStage {
+      if isSemaStage, let semanticTokens = semanticTokens(of: dict, for: snapshot) {
         // Only update semantic tokens if the 0,0 replacetext request returned semantic information.
-        updateSemanticTokens(response: dict, for: snapshot)
+        await semanticTokensManager.setSemanticTokens(for: snapshot.id, semanticTokens: semanticTokens)
       }
       if enablePublishDiagnostics {
         await publishDiagnostics(response: dict, for: snapshot, compileCommand: compileCommand)
@@ -494,7 +479,7 @@ extension SwiftLanguageServer {
     await self.publishDiagnostics(response: dict, for: snapshot, compileCommand: compileCommand)
   }
 
-  public func closeDocument(_ note: DidCloseTextDocumentNotification) {
+  public func closeDocument(_ note: DidCloseTextDocumentNotification) async {
     let keys = self.keys
 
     self.documentManager.close(note)
@@ -509,6 +494,8 @@ extension SwiftLanguageServer {
     self.currentDiagnostics[uri] = nil
 
     _ = try? self.sourcekitd.sendSync(req)
+
+    await semanticTokensManager.discaredSemanticTokens(for: note.textDocument.uri)
   }
 
   public func changeDocument(_ note: DidChangeTextDocumentNotification) async {
@@ -552,7 +539,16 @@ extension SwiftLanguageServer {
     guard let (preEditSnapshot, postEditSnapshot) = editResult else {
       return
     }
-    await syntaxTreeManager.registerEdit(preEditSnapshot: preEditSnapshot, postEditSnapshot: postEditSnapshot, edits: ConcurrentEdits(fromSequential: edits))
+    await syntaxTreeManager.registerEdit(
+      preEditSnapshot: preEditSnapshot,
+      postEditSnapshot: postEditSnapshot,
+      edits: ConcurrentEdits(fromSequential: edits)
+    )
+    await semanticTokensManager.registerEdit(
+      preEditSnapshot: preEditSnapshot.id,
+      postEditSnapshot: postEditSnapshot.id,
+      edits: note.contentChanges
+    )
 
     if let dict = lastResponse {
       let compileCommand = await self.buildSettings(for: note.textDocument.uri)