Remove document manager in SwiftLanguageService

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -17,6 +17,7 @@ import LanguageServerProtocolJSONRPC
 import SKCore
 import SKSupport
 import SwiftExtensions
+import SwiftSyntax
 
 import struct TSCBasic.AbsolutePath
 
@@ -443,7 +444,7 @@ extension ClangLanguageService {
 
   // MARK: - Text synchronization
 
-  public func openDocument(_ notification: DidOpenTextDocumentNotification) async {
+  public func openDocument(_ notification: DidOpenTextDocumentNotification, snapshot: DocumentSnapshot) async {
     openDocuments[notification.textDocument.uri] = notification.textDocument.language
     // Send clangd the build settings for the new file. We need to do this before
     // sending the open notification, so that the initial diagnostics already
@@ -459,7 +460,12 @@ extension ClangLanguageService {
 
   func reopenDocument(_ notification: ReopenTextDocumentNotification) {}
 
-  public func changeDocument(_ notification: DidChangeTextDocumentNotification) {
+  public func changeDocument(
+    _ notification: DidChangeTextDocumentNotification,
+    preEditSnapshot: DocumentSnapshot,
+    postEditSnapshot: DocumentSnapshot,
+    edits: [SourceEdit]
+  ) {
     clangd.send(notification)
   }
 

Sources/SourceKitLSP/DocumentManager.swift

@@ -204,42 +204,6 @@ public final class DocumentManager: InMemoryDocumentManager, Sendable {
   }
 }
 
-extension DocumentManager {
-
-  // MARK: - LSP notification handling
-
-  /// Convenience wrapper for `open(_:language:version:text:)` that logs on failure.
-  @discardableResult
-  func open(_ notification: DidOpenTextDocumentNotification) -> DocumentSnapshot? {
-    let doc = notification.textDocument
-    return orLog("failed to open document", level: .error) {
-      try open(doc.uri, language: doc.language, version: doc.version, text: doc.text)
-    }
-  }
-
-  /// Convenience wrapper for `close(_:)` that logs on failure.
-  func close(_ notification: DidCloseTextDocumentNotification) {
-    orLog("failed to close document", level: .error) {
-      try close(notification.textDocument.uri)
-    }
-  }
-
-  /// Convenience wrapper for `edit(_:newVersion:edits:updateDocumentTokens:)`
-  /// that logs on failure.
-  @discardableResult
-  func edit(
-    _ notification: DidChangeTextDocumentNotification
-  ) -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot, edits: [SourceEdit])? {
-    return orLog("failed to edit document", level: .error) {
-      return try edit(
-        notification.textDocument.uri,
-        newVersion: notification.textDocument.version,
-        edits: notification.contentChanges
-      )
-    }
-  }
-}
-
 fileprivate extension SourceEdit {
   /// Constructs a `SourceEdit` from the given `TextDocumentContentChangeEvent`.
   ///

Sources/SourceKitLSP/LanguageService.swift

@@ -13,6 +13,7 @@
 import Foundation
 import LanguageServerProtocol
 import SKCore
+import SwiftSyntax
 
 /// The state of a `ToolchainLanguageServer`
 public enum LanguageServerState {
@@ -124,9 +125,9 @@ public protocol LanguageService: AnyObject, Sendable {
   // MARK: - Text synchronization
 
   /// Sent to open up a document on the Language Server.
-  /// This may be called before or after a corresponding
-  /// `documentUpdatedBuildSettings` call for the same document.
-  func openDocument(_ notification: DidOpenTextDocumentNotification) async
+  ///
+  /// This may be called before or after a corresponding `documentUpdatedBuildSettings` call for the same document.
+  func openDocument(_ notification: DidOpenTextDocumentNotification, snapshot: DocumentSnapshot) async
 
   /// Sent to close a document on the Language Server.
   func closeDocument(_ notification: DidCloseTextDocumentNotification) async
@@ -138,7 +139,12 @@ public protocol LanguageService: AnyObject, Sendable {
   /// Only intended for `SwiftLanguageService`.
   func reopenDocument(_ notification: ReopenTextDocumentNotification) async
 
-  func changeDocument(_ notification: DidChangeTextDocumentNotification) async
+  func changeDocument(
+    _ notification: DidChangeTextDocumentNotification,
+    preEditSnapshot: DocumentSnapshot,
+    postEditSnapshot: DocumentSnapshot,
+    edits: [SourceEdit]
+  ) async
   func willSaveDocument(_ notification: WillSaveTextDocumentNotification) async
   func didSaveDocument(_ notification: DidSaveTextDocumentNotification) async
 

Sources/SourceKitLSP/Rename.swift

@@ -345,7 +345,7 @@ extension SwiftLanguageService {
     in uri: DocumentURI,
     name: CompoundDeclName
   ) async throws -> String {
-    guard let snapshot = documentManager.latestSnapshotOrDisk(uri, language: .swift) else {
+    guard let snapshot = try? documentManager.latestSnapshotOrDisk(uri, language: .swift) else {
       throw ResponseError.unknown("Failed to get contents of \(uri.forLogging) to translate Swift name to clang name")
     }
 

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1272,7 +1272,18 @@ extension SourceKitLSPServer {
   private func openDocument(_ notification: DidOpenTextDocumentNotification, workspace: Workspace) async {
     // Immediately open the document even if the build system isn't ready. This is important since
     // we check that the document is open when we receive messages from the build system.
-    documentManager.open(notification)
+    let snapshot = orLog("Opening document") {
+      try documentManager.open(
+        notification.textDocument.uri,
+        language: notification.textDocument.language,
+        version: notification.textDocument.version,
+        text: notification.textDocument.text
+      )
+    }
+    guard let snapshot else {
+      // Already logged failure
+      return
+    }
 
     let textDocument = notification.textDocument
     let uri = textDocument.uri
@@ -1286,7 +1297,7 @@ extension SourceKitLSPServer {
     await workspace.buildSystemManager.registerForChangeNotifications(for: uri, language: language)
 
     // If the document is ready, we can immediately send the notification.
-    await service.openDocument(notification)
+    await service.openDocument(notification, snapshot: snapshot)
   }
 
   func closeDocument(_ notification: DidCloseTextDocumentNotification) async {
@@ -1314,7 +1325,9 @@ extension SourceKitLSPServer {
   func closeDocument(_ notification: DidCloseTextDocumentNotification, workspace: Workspace) async {
     // Immediately close the document. We need to be sure to clear our pending work queue in case
     // the build system still isn't ready.
-    documentManager.close(notification)
+    orLog("failed to close document", level: .error) {
+      try documentManager.close(notification.textDocument.uri)
+    }
 
     let uri = notification.textDocument.uri
 
@@ -1335,8 +1348,23 @@ extension SourceKitLSPServer {
     await workspace.semanticIndexManager?.schedulePreparationForEditorFunctionality(of: uri)
 
     // If the document is ready, we can handle the change right now.
-    documentManager.edit(notification)
-    await workspace.documentService.value[uri]?.changeDocument(notification)
+    let editResult = orLog("Editing document") {
+      try documentManager.edit(
+        notification.textDocument.uri,
+        newVersion: notification.textDocument.version,
+        edits: notification.contentChanges
+      )
+    }
+    guard let (preEditSnapshot, postEditSnapshot, edits) = editResult else {
+      // Already logged failure
+      return
+    }
+    await workspace.documentService.value[uri]?.changeDocument(
+      notification,
+      preEditSnapshot: preEditSnapshot,
+      postEditSnapshot: postEditSnapshot,
+      edits: edits
+    )
   }
 
   func willSaveDocument(

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -93,7 +93,7 @@ public struct SwiftCompileCommand: Sendable, Equatable {
 }
 
 public actor SwiftLanguageService: LanguageService, Sendable {
-  /// The ``SourceKitLSPServer`` instance that created this `ClangLanguageService`.
+  /// The ``SourceKitLSPServer`` instance that created this `SwiftLanguageService`.
   weak var sourceKitLSPServer: SourceKitLSPServer?
 
   let sourcekitd: SourceKitD
@@ -121,9 +121,6 @@ public actor SwiftLanguageService: LanguageService, Sendable {
     options.generatedFilesAbsolutePath.asURL.appendingPathComponent("GeneratedMacroExpansions")
   }
 
-  // FIXME: ideally we wouldn't need separate management from a parent server in the same process.
-  var documentManager: DocumentManager
-
   /// For each edited document, the last task that was triggered to send a `PublishDiagnosticsNotification`.
   ///
   /// This is used to cancel previous publish diagnostics tasks if an edit is made to a document.
@@ -148,6 +145,15 @@ public actor SwiftLanguageService: LanguageService, Sendable {
 
   private var diagnosticReportManager: DiagnosticReportManager!
 
+  var documentManager: DocumentManager {
+    get throws {
+      guard let sourceKitLSPServer = self.sourceKitLSPServer else {
+        throw ResponseError.unknown("Connection to the editor closed")
+      }
+      return sourceKitLSPServer.documentManager
+    }
+  }
+
   /// Calling `scheduleCall` on `refreshDiagnosticsDebouncer` schedules a `DiagnosticsRefreshRequest` to be sent to
   /// to the client.
   ///
@@ -183,7 +189,6 @@ public actor SwiftLanguageService: LanguageService, Sendable {
     self.capabilityRegistry = workspace.capabilityRegistry
     self.semanticIndexManager = workspace.semanticIndexManager
     self.testHooks = testHooks
-    self.documentManager = DocumentManager()
     self.state = .connected
     self.diagnosticReportManager = nil  // Needed to work around rdar://116221716
     self.options = options
@@ -203,7 +208,7 @@ public actor SwiftLanguageService: LanguageService, Sendable {
       sourcekitd: self.sourcekitd,
       options: options,
       syntaxTreeManager: syntaxTreeManager,
-      documentManager: documentManager,
+      documentManager: try documentManager,
       clientHasDiagnosticsCodeDescriptionSupport: await self.clientHasDiagnosticsCodeDescriptionSupport
     )
 
@@ -439,15 +444,10 @@ extension SwiftLanguageService {
     ])
   }
 
-  public func openDocument(_ notification: DidOpenTextDocumentNotification) async {
+  public func openDocument(_ notification: DidOpenTextDocumentNotification, snapshot: DocumentSnapshot) async {
     cancelInFlightPublishDiagnosticsTask(for: notification.textDocument.uri)
     await diagnosticReportManager.removeItemsFromCache(with: notification.textDocument.uri)
 
-    guard let snapshot = self.documentManager.open(notification) else {
-      // Already logged failure.
-      return
-    }
-
     let buildSettings = await self.buildSettings(for: snapshot.uri)
 
     let req = openDocumentSourcekitdRequest(snapshot: snapshot, compileCommand: buildSettings)
@@ -460,8 +460,6 @@ extension SwiftLanguageService {
     inFlightPublishDiagnosticsTasks[notification.textDocument.uri] = nil
     await diagnosticReportManager.removeItemsFromCache(with: notification.textDocument.uri)
 
-    self.documentManager.close(notification)
-
     let req = closeDocumentSourcekitdRequest(uri: notification.textDocument.uri)
     _ = try? await self.sendSourcekitdRequest(req, fileContents: nil)
   }
@@ -542,7 +540,12 @@ extension SwiftLanguageService {
     }
   }
 
-  public func changeDocument(_ notification: DidChangeTextDocumentNotification) async {
+  public func changeDocument(
+    _ notification: DidChangeTextDocumentNotification,
+    preEditSnapshot: DocumentSnapshot,
+    postEditSnapshot: DocumentSnapshot,
+    edits: [SourceEdit]
+  ) async {
     cancelInFlightPublishDiagnosticsTask(for: notification.textDocument.uri)
 
     let keys = self.keys
@@ -552,10 +555,6 @@ extension SwiftLanguageService {
       let replacement: String
     }
 
-    guard let (preEditSnapshot, postEditSnapshot, edits) = self.documentManager.edit(notification) else {
-      return
-    }
-
     for edit in edits {
       let req = sourcekitd.dictionary([
         keys.request: self.requests.editorReplaceText,
@@ -1006,10 +1005,6 @@ extension SwiftLanguageService: SKDNotificationHandler {
 
     if notification.error == .connectionInterrupted {
       await self.setState(.connectionInterrupted)
-
-      // We don't have any open documents anymore after sourcekitd crashed.
-      // Reset the document manager to reflect that.
-      self.documentManager = DocumentManager()
     }
   }
 }