Merge pull request #2223 from ahoppen/neovim-tokens-out-of-sync

Fix issue that cause semantic tokens to be out-of-sync in neovim after a document changes on disk

Author: ahoppen

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -571,6 +571,7 @@ extension SwiftLanguageService {
     inFlightPublishDiagnosticsTasks[notification.textDocument.uri] = nil
     await diagnosticReportManager.removeItemsFromCache(with: notification.textDocument.uri)
     buildSettingsForOpenFiles[notification.textDocument.uri] = nil
+    await syntaxTreeManager.clearSyntaxTrees(for: notification.textDocument.uri)
     switch try? ReferenceDocumentURL(from: notification.textDocument.uri) {
     case .macroExpansion:
       break

Sources/SourceKitLSP/Swift/SyntaxTreeManager.swift

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import LanguageServerProtocol
 import SKUtilities
 import SwiftParser
 import SwiftSyntax
@@ -94,4 +95,9 @@ actor SyntaxTreeManager {
     }
     self.setComputation(for: postEditSnapshot.id, computation: incrementalParseComputation)
   }
+
+  /// Remove all cached syntax trees for the given document, eg. when the document is closed.
+  func clearSyntaxTrees(for uri: DocumentURI) {
+    syntaxTreeComputations.removeAll(where: { $0.uri == uri })
+  }
 }

Tests/SourceKitLSPTests/SemanticTokensTests.swift

@@ -913,6 +913,55 @@ final class SemanticTokensTests: XCTestCase {
     )
   }
 
+  func testCloseAndReopenDocumentWithSameDocumentVersion() async throws {
+    // When neovim detects a change of the document on-disk (eg. caused by git operations). It closes the document and
+    // re-opens it with the same document version but different contents. Check that we don't re-use the syntax tree of
+    // the previously opened document.
+    let testClient = try await TestSourceKitLSPClient()
+    let uri = DocumentURI(for: .swift)
+    let initialPositions = testClient.openDocument(
+      """
+      1️⃣import 2️⃣Foo
+      3️⃣func 4️⃣bar() {}
+      """,
+      uri: uri
+    )
+    let initialTokens = try await testClient.send(
+      DocumentSemanticTokensRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    XCTAssertEqual(
+      SyntaxHighlightingTokens(lspEncodedTokens: try unwrap(initialTokens).data).tokens,
+      [
+        Token(start: initialPositions["1️⃣"], utf16length: 6, kind: .keyword),
+        Token(start: initialPositions["2️⃣"], utf16length: 3, kind: .identifier),
+        Token(start: initialPositions["3️⃣"], utf16length: 4, kind: .keyword),
+        Token(start: initialPositions["4️⃣"], utf16length: 3, kind: .identifier),
+      ]
+    )
+
+    testClient.send(DidCloseTextDocumentNotification(textDocument: TextDocumentIdentifier(uri)))
+
+    let reopenedPositions = testClient.openDocument(
+      """
+      1️⃣func 2️⃣bar() {}
+      """,
+      uri: uri
+    )
+
+    let reopenedTokens = try await testClient.send(
+      DocumentSemanticTokensRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    XCTAssertEqual(
+      SyntaxHighlightingTokens(lspEncodedTokens: try unwrap(reopenedTokens).data).tokens,
+      [
+        Token(start: reopenedPositions["1️⃣"], utf16length: 4, kind: .keyword),
+        Token(start: reopenedPositions["2️⃣"], utf16length: 3, kind: .identifier),
+      ]
+    )
+  }
+
   func testClang() async throws {
     try await assertSemanticTokens(
       markedContents: """