Store syntax trees in semantic token in syntax tree/semantic token managers instead of in the Document

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
@@ -25,11 +24,14 @@ target_sources(SourceKitLSP PRIVATE
   Swift/OpenInterface.swift
   Swift/SemanticRefactorCommand.swift
   Swift/SemanticRefactoring.swift
+  Swift/SemanticTokens.swift
+  Swift/SemanticTokensManager.swift
   Swift/SourceKitD+ResponseError.swift
   Swift/SwiftCommand.swift
   Swift/SwiftLanguageServer.swift
   Swift/SyntaxHighlightingToken.swift
   Swift/SyntaxHighlightingTokenParser.swift
+  Swift/SyntaxTreeManager.swift
   Swift/VariableTypeInfo.swift)
 set_target_properties(SourceKitLSP PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})

Sources/SourceKitLSP/DocumentManager.swift

@@ -15,27 +15,46 @@ import LanguageServerProtocol
 import LSPLogging
 import SKSupport
 
-public struct DocumentSnapshot {
-  public var document: Document
-  public var version: Int
-  public var 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 var tokens: DocumentTokens
+/// An immutable snapshot of a document at a given time.
+///
+/// ``DocumentSnapshot`` is always derived from a ``Document``. That is, the
+/// data structure that is stored internally by the ``DocumentManager`` is a
+/// ``Document``. The purpose of a ``DocumentSnapshot`` is to be able to work
+/// with one version of a document without having to think about it changing.
+public struct DocumentSnapshot: Identifiable {
+  /// An ID that uniquely identifies the version of the document stored in this
+  /// snapshot.
+  public struct ID: Hashable, Comparable {
+    public let uri: DocumentURI
+    public let version: Int
+
+    /// Returns `true` if the snapshots reference the same document but rhs has a
+    /// later version than `lhs`.
+    ///
+    /// Snapshot IDs of different documents are not comparable to each other and
+    /// will always return `false`.
+    public static func < (lhs: DocumentSnapshot.ID, rhs: DocumentSnapshot.ID) -> Bool {
+      return lhs.uri == rhs.uri && lhs.version < rhs.version
+    }
+  }
+
+  public let id: ID
+  public let language: Language
+  public let lineTable: LineTable
 
+  public var uri: DocumentURI { id.uri }
+  public var version: Int { id.version }
   public var text: String { lineTable.content }
 
   public init(
-    document: Document,
+    uri: DocumentURI,
+    language: Language,
     version: Int,
-    lineTable: LineTable,
-    tokens: DocumentTokens
+    lineTable: LineTable
   ) {
-    self.document = document
-    self.version = version
+    self.id = ID(uri: uri, version: version)
+    self.language = language
     self.lineTable = lineTable
-    self.tokens = tokens
   }
 
   func index(of pos: Position) -> String.Index? {
@@ -48,23 +67,21 @@ 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.
   fileprivate var latestSnapshot: DocumentSnapshot {
     DocumentSnapshot(
-      document: self,
+      uri: self.uri,
+      language: self.language,
       version: latestVersion,
-      lineTable: latestLineTable,
-      tokens: latestTokens
+      lineTable: latestLineTable
     )
   }
 }
@@ -118,28 +135,31 @@ public final class DocumentManager {
 
   /// Applies the given edits to the document.
   ///
-  /// - parameter willEditDocument: Optional closure to call before each edit.
-  /// - parameter updateDocumentTokens: 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.
+  /// - Parameters:
+  ///   - uri: The URI of the document to update
+  ///   - newVersion: The new version of the document. Must be greater than the
+  ///     latest version of the document.
+  ///   - edits: The edits to apply to the document
+  ///   - willEditDocument: Optional closure to call before each edit. Will be 
+  ///     called multiple times if there are multiple edits.
+  /// - Returns: The snapshot of the document before the edit and the snapshot
+  ///   of the document after the edit.
   @discardableResult
   public func edit(
     _ uri: DocumentURI,
     newVersion: Int,
     edits: [TextDocumentContentChangeEvent],
-    willEditDocument: ((_ before: DocumentSnapshot, TextDocumentContentChangeEvent) -> Void)? = nil,
-    updateDocumentTokens: ((_ after: DocumentSnapshot) -> DocumentTokens)? = nil
-  ) throws -> DocumentSnapshot {
+    willEditDocument: ((_ before: LineTable, TextDocumentContentChangeEvent) -> Void)? = nil
+  ) throws -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot) {
     return try queue.sync {
       guard let document = documents[uri] else {
         throw Error.missingDocument(uri)
       }
+      let preEditSnapshot = document.latestSnapshot
 
       for edit in edits {
-        if let f = willEditDocument {
-          f(document.latestSnapshot, edit)
+        if let willEditDocument {
+          willEditDocument(document.latestLineTable, edit)
         }
 
         if let range = edit.range  {
@@ -149,49 +169,17 @@ 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()
-        }
-
-        if let f = updateDocumentTokens {
-          document.latestTokens = f(document.latestSnapshot)
         }
       }
 
-      document.latestVersion = newVersion
-      return document.latestSnapshot
-    }
-  }
-
-  /// 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)
+      if newVersion <= document.latestVersion {
+        log("Document version did not increase on edit from \(document.latestVersion) to \(newVersion)", level: .error)
       }
-
-      document.latestTokens = tokens
-
-      return document.latestSnapshot
+      document.latestVersion = newVersion
+      return (preEditSnapshot, document.latestSnapshot)
     }
   }
 
@@ -230,16 +218,14 @@ extension DocumentManager {
   @discardableResult
   func edit(
     _ note: DidChangeTextDocumentNotification,
-    willEditDocument: ((_ before: DocumentSnapshot, TextDocumentContentChangeEvent) -> Void)? = nil,
-    updateDocumentTokens: ((_ after: DocumentSnapshot) -> DocumentTokens)? = nil
-  ) -> DocumentSnapshot? {
+    willEditDocument: ((_ before: LineTable, TextDocumentContentChangeEvent) -> Void)? = nil
+  ) -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot)? {
     return orLog("failed to edit document", level: .error) {
-      try edit(
+      return try edit(
         note.textDocument.uri,
         newVersion: note.textDocument.version,
         edits: note.contentChanges,
-        willEditDocument: willEditDocument,
-        updateDocumentTokens: updateDocumentTokens
+        willEditDocument: willEditDocument
       )
     }
   }

Sources/SourceKitLSP/SourceKitServer.swift

@@ -374,7 +374,7 @@ public actor SourceKitServer {
       await self.closeDocument(closeNotification, workspace: workspace)
 
       let textDocument = TextDocumentItem(uri: documentUri,
-                                          language: snapshot.document.language,
+                                          language: snapshot.language,
                                           version: snapshot.version,
                                           text: snapshot.text)
       await self.openDocument(DidOpenTextDocumentNotification(textDocument: textDocument), workspace: workspace)
@@ -1146,7 +1146,7 @@ extension SourceKitServer {
         if let newWorkspace = newWorkspace {
           await self.openDocument(DidOpenTextDocumentNotification(textDocument: TextDocumentItem(
             uri: docUri,
-            language: snapshot.document.language,
+            language: snapshot.language,
             version: snapshot.version,
             text: snapshot.text
           )), workspace: newWorkspace)

Sources/SourceKitLSP/Swift/CodeCompletion.swift

@@ -63,8 +63,8 @@ extension SwiftLanguageServer {
         log("triggerFromIncompleteCompletions with no existing completion session", level: .warning)
         throw ResponseError.serverCancelled
       }
-      guard currentSession.uri == snapshot.document.uri, currentSession.utf8StartOffset == offset else {
-        log("triggerFromIncompleteCompletions with incompatible completion session; expected \(currentSession.uri)@\(currentSession.utf8StartOffset), but got \(snapshot.document.uri)@\(offset)", level: .warning)
+      guard currentSession.uri == snapshot.uri, currentSession.utf8StartOffset == offset else {
+        log("triggerFromIncompleteCompletions with incompatible completion session; expected \(currentSession.uri)@\(currentSession.utf8StartOffset), but got \(snapshot.uri)@\(offset)", level: .warning)
         throw ResponseError.serverCancelled
       }
       session = currentSession
@@ -76,7 +76,7 @@ extension SwiftLanguageServer {
         snapshot: snapshot,
         utf8Offset: offset,
         position: completionPos,
-        compileCommand: await buildSettings(for: snapshot.document.uri))
+        compileCommand: await buildSettings(for: snapshot.uri))
 
       await currentCompletionSession?.close()
       currentCompletionSession = session
@@ -95,15 +95,15 @@ extension SwiftLanguageServer {
     let skreq = SKDRequestDictionary(sourcekitd: sourcekitd)
     skreq[keys.request] = requests.codecomplete
     skreq[keys.offset] = offset
-    skreq[keys.sourcefile] = snapshot.document.uri.pseudoPath
+    skreq[keys.sourcefile] = snapshot.uri.pseudoPath
     skreq[keys.sourcetext] = snapshot.text
 
     let skreqOptions = SKDRequestDictionary(sourcekitd: sourcekitd)
     skreqOptions[keys.codecomplete_sort_byname] = 1
     skreq[keys.codecomplete_options] = skreqOptions
 
     // FIXME: SourceKit should probably cache this for us.
-    if let compileCommand = await buildSettings(for: snapshot.document.uri) {
+    if let compileCommand = await buildSettings(for: snapshot.uri) {
       skreq[keys.compilerargs] = compileCommand.compilerArgs
     }
 

Sources/SourceKitLSP/Swift/CodeCompletionSession.swift

@@ -41,7 +41,7 @@ actor CodeCompletionSession {
     case open
   }
 
-  nonisolated var uri: DocumentURI { snapshot.document.uri }
+  nonisolated var uri: DocumentURI { snapshot.uri }
 
   init(
     server: SwiftLanguageServer,
@@ -174,7 +174,7 @@ actor CodeCompletionSession {
     let keys = server.sourcekitd.keys
     req[keys.request] = server.sourcekitd.requests.codecomplete_close
     req[keys.offset] = self.utf8StartOffset
-    req[keys.name] = self.snapshot.document.uri.pseudoPath
+    req[keys.name] = self.snapshot.uri.pseudoPath
     log("\(Self.self) Closing: \(self)")
     _ = try? server.sourcekitd.sendSync(req)
   }

Sources/SourceKitLSP/Swift/CursorInfo.swift

@@ -101,7 +101,7 @@ extension SwiftLanguageServer {
     if offsetRange.upperBound != offsetRange.lowerBound {
       skreq[keys.length] = offsetRange.count
     }
-    skreq[keys.sourcefile] = snapshot.document.uri.pseudoPath
+    skreq[keys.sourcefile] = snapshot.uri.pseudoPath
 
     // FIXME: SourceKit should probably cache this for us.
     if let compileCommand = await self.buildSettings(for: uri) {

Sources/SourceKitLSP/Swift/Diagnostic.swift

@@ -65,7 +65,7 @@ extension CodeAction {
       title: title,
       kind: .quickFix,
       diagnostics: nil,
-      edit: WorkspaceEdit(changes: [snapshot.document.uri:edits]))
+      edit: WorkspaceEdit(changes: [snapshot.uri:edits]))
   }
 
   /// Describe a fixit's edit briefly.
@@ -257,7 +257,7 @@ extension DiagnosticRelatedInformation {
     }
 
     self.init(
-      location: Location(uri: snapshot.document.uri, range: Range(position!)),
+      location: Location(uri: snapshot.uri, range: Range(position!)),
       message: message,
       codeActions: actions)
   }

Sources/SourceKitLSP/Swift/ExpressionTypeInfo.swift

@@ -66,7 +66,7 @@ extension SwiftLanguageServer {
 
     let skreq = SKDRequestDictionary(sourcekitd: sourcekitd)
     skreq[keys.request] = requests.expression_type
-    skreq[keys.sourcefile] = snapshot.document.uri.pseudoPath
+    skreq[keys.sourcefile] = snapshot.uri.pseudoPath
 
     // FIXME: SourceKit should probably cache this for us.
     if let compileCommand = await self.buildSettings(for: uri) {

Sources/SourceKitLSP/Swift/SemanticRefactoring.swift

@@ -74,7 +74,7 @@ struct SemanticRefactoring {
     }
 
     self.title = title
-    self.edit = WorkspaceEdit(changes: [snapshot.document.uri: textEdits])
+    self.edit = WorkspaceEdit(changes: [snapshot.uri: textEdits])
   }
 }
 
@@ -155,7 +155,7 @@ extension SwiftLanguageServer {
     skreq[keys.actionuid] = self.sourcekitd.api.uid_get_from_cstr(refactorCommand.actionString)!
 
     // FIXME: SourceKit should probably cache this for us.
-    if let compileCommand = await self.buildSettings(for: snapshot.document.uri) {
+    if let compileCommand = await self.buildSettings(for: snapshot.uri) {
       skreq[keys.compilerargs] = compileCommand.compilerArgs
     }