Refactor logic for updating lexical/syntactic tokens

- Fix typo
- Refactor update tokens logic
  Remove the replaceXTokens methods from DocumentManager and implement
  updates as changes on DocumentTokens directly.
- Add after-edit callback and use it to update tokens

Author: fwcd

Sources/SourceKitLSP/CapabilityRegistry.swift

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

Sources/SourceKitLSP/DocumentManager.swift

@@ -37,6 +37,12 @@ public struct DocumentTokens {
     action(&syntactic)
     action(&semantic)
   }
+
+  // Replace all lexical tokens in `range`.
+  public mutating func replaceLexical(in range: Range<Position>, with newTokens: [SyntaxHighlightingToken]) {
+    lexical.removeAll { $0.range.overlaps(range) }
+    lexical += newTokens
+  }
 }
 
 public struct DocumentSnapshot {
@@ -140,19 +146,27 @@ public final class DocumentManager {
 
   /// Applies the given edits to the document.
   ///
-  /// - parameter editCallback: Optional closure to call for each edit.
+  /// - parameter beforeCallback: Optional closure to call before each edit.
+  /// - parameter afterCallback: Optional closure to call after each edit.
   /// - parameter before: The document contents *before* the edit is applied.
+  /// - parameter after: The document contents *after* the edit is applied.
   /// - returns: The contents of the file after all the edits are applied.
   /// - throws: Error.missingDocument if the document is not open.
   @discardableResult
-  public func edit(_ uri: DocumentURI, newVersion: Int, edits: [TextDocumentContentChangeEvent], editCallback: ((_ before: DocumentSnapshot, TextDocumentContentChangeEvent) -> Void)? = nil) throws -> DocumentSnapshot {
+  public func edit(
+    _ uri: DocumentURI,
+    newVersion: Int,
+    edits: [TextDocumentContentChangeEvent],
+    beforeCallback: ((_ before: DocumentSnapshot, TextDocumentContentChangeEvent) -> Void)? = nil,
+    afterCallback: ((_ after: DocumentSnapshot) -> DocumentTokens?)? = nil
+  ) throws -> DocumentSnapshot {
     return try queue.sync {
       guard let document = documents[uri] else {
         throw Error.missingDocument(uri)
       }
 
       for edit in edits {
-        if let f = editCallback {
+        if let f = beforeCallback {
           f(document.latestSnapshot, edit)
         }
 
@@ -199,75 +213,28 @@ public final class DocumentManager {
           document.latestTokens = DocumentTokens()
         }
 
+        if let f = afterCallback, let tokens = f(document.latestSnapshot) {
+          document.latestTokens = tokens
+        }
       }
 
       document.latestVersion = newVersion
       return document.latestSnapshot
     }
   }
 
-  /// Replaces the semantic tokens for a document.
+  /// Updates the tokens in a document.
   ///
   /// - parameter uri: The URI of the document to be updated
-  /// - parameter tokens: The tokens to be used
+  /// - parameter tokens: The new tokens for the document
   @discardableResult
-  public func replaceSemanticTokens(
-    _ uri: DocumentURI,
-    tokens: [SyntaxHighlightingToken]
-  ) throws -> DocumentSnapshot {
+  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.semantic = tokens
-      return document.latestSnapshot
-    }
-  }
-
-  /// Replaces the syntactic tokens for a document.
-  ///
-  /// - parameter uri: The URI of the document to be updated
-  /// - parameter tokens: The tokens to be used
-  @discardableResult
-  public func replaceSyntacticTokens(
-    _ uri: DocumentURI,
-    tokens: [SyntaxHighlightingToken]
-  ) throws -> DocumentSnapshot {
-    return try queue.sync {
-      guard let document = documents[uri] else {
-        throw Error.missingDocument(uri)
-      }
-
-      document.latestTokens.syntactic = tokens
-      return document.latestSnapshot
-    }
-  }
-
-  /// Replaces the given lexical tokens in a document
-  /// within the given range.
-  ///
-  /// - parameter uri: The URI of the document to be updated
-  /// - parameter range: The range to replace tokens in (nil means the entire document)
-  /// - parameter newTokens: The tokens to be added
-  @discardableResult
-  public func replaceLexicalTokens(
-    _ uri: DocumentURI,
-    in range: Range<Position>? = nil,
-    with newTokens: [SyntaxHighlightingToken]
-  ) throws -> DocumentSnapshot {
-    return try queue.sync {
-      guard let document = documents[uri] else {
-        throw Error.missingDocument(uri)
-      }
-
-      // Remove all tokens that overlap with `range`
-      // (or the entire document if `range` is `nil`)
-      document.latestTokens.lexical.removeAll { token in
-        range.map { token.range.overlaps($0) } ?? true
-      }
-
-      document.latestTokens.lexical += newTokens
+      document.latestTokens = tokens
 
       return document.latestSnapshot
     }
@@ -303,11 +270,21 @@ extension DocumentManager {
     }
   }
 
-  /// Convenience wrapper for `edit(_:newVersion:edits:editCallback:)` that logs on failure.
+  /// Convenience wrapper for `edit(_:newVersion:edits:beforeCallback:)` that logs on failure.
   @discardableResult
-  func edit(_ note: DidChangeTextDocumentNotification, editCallback: ((_ before: DocumentSnapshot, TextDocumentContentChangeEvent) -> Void)? = nil) -> DocumentSnapshot? {
+  func edit(
+    _ note: DidChangeTextDocumentNotification,
+    beforeCallback: ((_ before: DocumentSnapshot, TextDocumentContentChangeEvent) -> Void)? = nil,
+    afterCallback: ((_ after: DocumentSnapshot) -> DocumentTokens?)? = nil
+  ) -> DocumentSnapshot? {
     return orLog("failed to edit document", level: .error) {
-      try edit(note.textDocument.uri, newVersion: note.textDocument.version ?? -1, edits: note.contentChanges, editCallback: editCallback)
+      try edit(
+        note.textDocument.uri,
+        newVersion: note.textDocument.version ?? -1,
+        edits: note.contentChanges,
+        beforeCallback: beforeCallback,
+        afterCallback: afterCallback
+      )
     }
   }
 }

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -161,77 +161,100 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
     }
   }
 
-  /// Update lexical and syntactic tokens for the given `snapshot`.
-  /// Should only be called on `queue`.
+  /// Updates the lexical and syntactic tokens for the given `snapshot`.
+  /// Must be called on `self.queue`.
   private func updateLexicalAndSyntacticTokens(
     response: SKDResponseDictionary,
     for snapshot: DocumentSnapshot
   ) {
     dispatchPrecondition(condition: .onQueue(queue))
 
+    let uri = snapshot.document.uri
+    let docTokens = updatedLexicalAndSyntacticTokens(response: response, for: snapshot)
+
+    do {
+      try documentManager.updateTokens(uri, tokens: docTokens)
+    } catch {
+      log("Updating lexical and syntactic tokens failed: \(error)", level: .warning)
+    }
+  }
+
+  /// Returns the updated lexical and syntactic tokens for the given `snapshot`.
+  private func updatedLexicalAndSyntacticTokens(
+    response: SKDResponseDictionary,
+    for snapshot: DocumentSnapshot
+  ) -> DocumentTokens {
+    var docTokens = snapshot.tokens
+
     guard let offset: Int = response[keys.offset],
           let length: Int = response[keys.length],
           let start: Position = snapshot.positionOf(utf8Offset: offset),
           let end: Position = snapshot.positionOf(utf8Offset: offset + length) else {
       log("updateLexicalAndSyntacticTokens failed, no range found", level: .error)
-      return
+      return docTokens
     }
 
-    let uri = snapshot.document.uri
     let range = start..<end
 
     // If the range is empty we don't have to (and shouldn't) update anything.
     // This is important, since the substructure may be empty, causing us to
     // unnecessarily remove all syntactic tokens.
     guard !range.isEmpty else {
-      return
+      return docTokens
     }
 
     if let syntaxMap: SKDResponseArray = response[keys.syntaxmap] {
       let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd)
       let tokens = tokenParser.parseTokens(syntaxMap, in: snapshot)
 
-      do {
-        try documentManager.replaceLexicalTokens(uri, in: range, with: tokens)
-      } catch {
-        log("updating lexical tokens for \(uri) failed: \(error)", level: .warning)
-      }
+      docTokens.replaceLexical(in: range, with: tokens)
     }
 
     if let substructure: SKDResponseArray = response[keys.substructure] {
       let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd, useName: true)
       let tokens = tokenParser.parseTokens(substructure, in: snapshot)
-      do {
-        try documentManager.replaceSyntacticTokens(uri, tokens: tokens)
-      } catch {
-        log("updating syntactic tokens for \(uri) failed: \(error)", level: .warning)
-      }
+
+      docTokens.syntactic = tokens
     }
+
+    return docTokens
   }
 
-  /// Update semantic tokens for the given `snapshot`.
-  /// Should only be called on `queue`.
+  /// Updates the semantic tokens for the given `snapshot`.
+  /// Must be called on `self.queue`.
   private func updateSemanticTokens(
     response: SKDResponseDictionary,
     for snapshot: DocumentSnapshot
   ) {
     dispatchPrecondition(condition: .onQueue(queue))
 
     let uri = snapshot.document.uri
-    guard let skTokens: SKDResponseArray = response[keys.annotations] else {
-      return
-    }
-
-    let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd)
-    let tokens = tokenParser.parseTokens(skTokens, in: snapshot)
+    let docTokens = updatedSemanticTokens(response: response, for: snapshot)
 
     do {
-      try documentManager.replaceSemanticTokens(uri, tokens: tokens)
+      try documentManager.updateTokens(uri, tokens: docTokens)
     } catch {
-      log("updating semantic tokens for \(uri) failed: \(error)", level: .warning)
+      log("Updating semantic tokens failed: \(error)", level: .warning)
     }
   }
 
+  /// Returns the updated semantic tokens for the given `snapshot`.
+  private func updatedSemanticTokens(
+    response: SKDResponseDictionary,
+    for snapshot: DocumentSnapshot
+  ) -> DocumentTokens {
+    var docTokens = snapshot.tokens
+
+    if let skTokens: SKDResponseArray = response[keys.annotations] {
+      let tokenParser = SyntaxHighlightingTokenParser(sourcekitd: sourcekitd)
+      let tokens = tokenParser.parseTokens(skTokens, in: snapshot)
+
+      docTokens.semantic = tokens
+    }
+
+    return docTokens
+  }
+
   /// Inform the client about changes to the syntax highlighting tokens.
   private func requestTokensRefresh() {
     if clientCapabilities.workspace?.semanticTokens?.refreshSupport ?? false {
@@ -491,50 +514,35 @@ extension SwiftLanguageServer {
 
     self.queue.async {
       let uri = note.textDocument.uri
-      let newVersion = note.textDocument.version ?? -1
       var lastResponse: SKDResponseDictionary? = nil
 
-      // We iterate manually over the `contentChanges` here instead of passing
-      // the notification directly to the convenience method `edit(_:editCallback:)`
-      // since we need to update the tokens in lockstep with the edits in the
-      // document.
-
-      do {
-        for edit in note.contentChanges {
-          try self.documentManager.edit(uri, newVersion: newVersion, edits: [edit]) { (before: DocumentSnapshot, edit: TextDocumentContentChangeEvent) in
-            let req = SKDRequestDictionary(sourcekitd: self.sourcekitd)
-            req[keys.request] = self.requests.editor_replacetext
-            req[keys.name] = note.textDocument.uri.pseudoPath
-
-            if let range = edit.range {
-              guard let offset = before.utf8Offset(of: range.lowerBound), let end = before.utf8Offset(of: range.upperBound) else {
-                fatalError("invalid edit \(range)")
-              }
+      self.documentManager.edit(note) { (before: DocumentSnapshot, edit: TextDocumentContentChangeEvent) in
+        let req = SKDRequestDictionary(sourcekitd: self.sourcekitd)
+        req[keys.request] = self.requests.editor_replacetext
+        req[keys.name] = note.textDocument.uri.pseudoPath
 
-              req[keys.offset] = offset
-              req[keys.length] = end - offset
-
-            } else {
-              // Full text
-              req[keys.offset] = 0
-              req[keys.length] = before.text.utf8.count
-            }
-
-            req[keys.sourcetext] = edit.text
-
-            lastResponse = try? self.sourcekitd.sendSync(req)
+        if let range = edit.range {
+          guard let offset = before.utf8Offset(of: range.lowerBound), let end = before.utf8Offset(of: range.upperBound) else {
+            fatalError("invalid edit \(range)")
           }
 
-          // SourceKit seems to respond with an empty substructure after sending an
-          // empty range for an edit, causing all syntactic tokens to get removed
-          // therefore we only update them if the range is non-empty.
+          req[keys.offset] = offset
+          req[keys.length] = end - offset
 
-          if let dict = lastResponse, let snapshot = self.documentManager.latestSnapshot(uri) {
-            self.updateLexicalAndSyntacticTokens(response: dict, for: snapshot)
-          }
+        } else {
+          // Full text
+          req[keys.offset] = 0
+          req[keys.length] = before.text.utf8.count
+        }
+
+        req[keys.sourcetext] = edit.text
+        lastResponse = try? self.sourcekitd.sendSync(req)
+      } afterCallback: { (after: DocumentSnapshot) in
+        if let dict = lastResponse {
+          return self.updatedLexicalAndSyntacticTokens(response: dict, for: after)
+        } else {
+          return nil
         }
-      } catch {
-        log("failed to edit document: \(error)", level: .error)
       }
 
       if let dict = lastResponse, let snapshot = self.documentManager.latestSnapshot(uri) {