[DocC Live Preview] Cache on-disk snapshots opened in sourcekitd

Sources/ClangLanguageService/ClangLanguageService.swift

@@ -421,6 +421,14 @@ extension ClangLanguageService {
     clangd.send(notification)
   }
 
+  package func openDocumentOnDisk(snapshot: DocumentSnapshot, originalFile: DocumentURI) async throws {
+    // clangd doesn't support on-disk documents
+  }
+
+  package func closeDocumentOnDisk(snapshot: DocumentSnapshot) async {
+    // clangd doesn't support on-disk documents
+  }
+
   package func reopenDocument(_ notification: ReopenTextDocumentNotification) {}
 
   package func changeDocument(
@@ -512,9 +520,9 @@ extension ClangLanguageService {
     return try await forwardRequestToClangd(req)
   }
 
-  package func symbolGraph(
-    forOnDiskContentsOf symbolDocumentUri: DocumentURI,
-    at location: SymbolLocation
+  package func symbolGraphForDocumentOnDisk(
+    at location: SymbolLocation,
+    manager: OnDiskDocumentManager
   ) async throws -> String? {
     return nil
   }

Sources/DocCDocumentation/IndexStoreDB+Extensions.swift

@@ -40,7 +40,7 @@ extension CheckedIndex {
     var result: [SymbolOccurrence] = []
     for occurrence in topLevelSymbolOccurrences {
       let info = try await doccSymbolInformation(ofUSR: occurrence.symbol.usr, fetchSymbolGraph: fetchSymbolGraph)
-      if let info, info.matches(symbolLink) {
+      if info.matches(symbolLink) {
         result.append(occurrence)
       }
     }
@@ -60,9 +60,9 @@ extension CheckedIndex {
   package func doccSymbolInformation(
     ofUSR usr: String,
     fetchSymbolGraph: (SymbolLocation) async throws -> String?
-  ) async throws -> DocCSymbolInformation? {
+  ) async throws -> DocCSymbolInformation {
     guard let topLevelSymbolOccurrence = primaryDefinitionOrDeclarationOccurrence(ofUSR: usr) else {
-      return nil
+      throw DocCCheckedIndexError.emptyDocCSymbolLink
     }
     let moduleName = topLevelSymbolOccurrence.location.moduleName
     var symbols = [topLevelSymbolOccurrence]

Sources/SourceKitLSP/CMakeLists.txt

@@ -13,6 +13,7 @@ add_library(SourceKitLSP STATIC
   LogMessageNotification+representingStructureUsingEmojiPrefixIfNecessary.swift
   MacroExpansionReferenceDocumentURLData.swift
   MessageHandlingDependencyTracker.swift
+  OnDiskDocumentManager.swift
   ReferenceDocumentURL.swift
   Rename.swift
   SemanticTokensLegend+SourceKitLSPLegend.swift
@@ -27,7 +28,7 @@ add_library(SourceKitLSP STATIC
   Workspace.swift
 )
 target_sources(SourceKitLSP PRIVATE
-  Documentation/DocCDocumentationHandler.swift
+  Documentation/DocCDocumentation.swift
   Documentation/DocumentationLanguageService.swift
 )
 set_target_properties(SourceKitLSP PROPERTIES

Sources/SourceKitLSP/Documentation/DoccDocumentationHandler.swift → Sources/SourceKitLSP/Documentation/DoccDocumentation.swift

@@ -58,42 +58,34 @@ extension DocumentationLanguageService {
         guard let index = workspace.index(checkedFor: .deletedFiles) else {
           throw ResponseError.requestFailed(doccDocumentationError: .indexNotAvailable)
         }
-        guard let symbolLink = DocCSymbolLink(linkString: symbolName),
-          let symbolOccurrence = try await index.primaryDefinitionOrDeclarationOccurrence(
-            ofDocCSymbolLink: symbolLink,
-            fetchSymbolGraph: { location in
-              guard let symbolWorkspace = try await workspaceForDocument(uri: location.documentUri),
-                let languageService = try await languageService(for: location.documentUri, .swift, in: symbolWorkspace)
-              else {
-                throw ResponseError.internalError("Unable to find Swift language service for \(location.documentUri)")
+        return try await sourceKitLSPServer.withOnDiskDocumentManager { onDiskDocumentManager in
+          guard let symbolLink = DocCSymbolLink(linkString: symbolName),
+            let symbolOccurrence = try await index.primaryDefinitionOrDeclarationOccurrence(
+              ofDocCSymbolLink: symbolLink,
+              fetchSymbolGraph: {
+                try await onDiskDocumentManager.languageService(for: $0.documentUri, .swift)
+                  .symbolGraphForDocumentOnDisk(at: $0, manager: onDiskDocumentManager)
               }
-              return try await languageService.symbolGraph(forOnDiskContentsOf: location.documentUri, at: location)
-            }
+            )
+          else {
+            throw ResponseError.requestFailed(doccDocumentationError: .symbolNotFound(symbolName))
+          }
+          guard
+            let symbolGraph = try await onDiskDocumentManager.languageService(
+              for: symbolOccurrence.location.documentUri,
+              .swift
+            ).symbolGraphForDocumentOnDisk(at: symbolOccurrence.location, manager: onDiskDocumentManager)
+          else {
+            throw ResponseError.internalError("Unable to retrieve symbol graph for \(symbolOccurrence.symbol.name)")
+          }
+          return try await documentationManager.renderDocCDocumentation(
+            symbolUSR: symbolOccurrence.symbol.usr,
+            symbolGraph: symbolGraph,
+            markupFile: snapshot.text,
+            moduleName: moduleName,
+            catalogURL: catalogURL
           )
-        else {
-          throw ResponseError.requestFailed(doccDocumentationError: .symbolNotFound(symbolName))
         }
-        let symbolDocumentUri = symbolOccurrence.location.documentUri
-        guard
-          let symbolWorkspace = try await workspaceForDocument(uri: symbolDocumentUri),
-          let languageService = try await languageService(for: symbolDocumentUri, .swift, in: symbolWorkspace)
-        else {
-          throw ResponseError.internalError("Unable to find Swift language service for \(symbolDocumentUri)")
-        }
-        let symbolGraph = try await languageService.symbolGraph(
-          forOnDiskContentsOf: symbolDocumentUri,
-          at: symbolOccurrence.location
-        )
-        guard let symbolGraph else {
-          throw ResponseError.internalError("Unable to retrieve symbol graph for \(symbolOccurrence.symbol.name)")
-        }
-        return try await documentationManager.renderDocCDocumentation(
-          symbolUSR: symbolOccurrence.symbol.usr,
-          symbolGraph: symbolGraph,
-          markupFile: snapshot.text,
-          moduleName: moduleName,
-          catalogURL: catalogURL
-        )
       }
       // This is a page representing the module itself.
       // Create a dummy symbol graph and tell SwiftDocC to convert the module name.

Sources/SourceKitLSP/Documentation/DocumentationLanguageService.swift

@@ -98,6 +98,14 @@ package actor DocumentationLanguageService: LanguageService, Sendable {
     // The DocumentationLanguageService does not do anything with document events
   }
 
+  package func openDocumentOnDisk(snapshot: DocumentSnapshot, originalFile: DocumentURI) async throws {
+    // The DocumentationLanguageService does not do anything with document events
+  }
+
+  package func closeDocumentOnDisk(snapshot: DocumentSnapshot) async {
+    // The DocumentationLanguageService does not do anything with document events
+  }
+
   package func reopenDocument(_ notification: ReopenTextDocumentNotification) async {
     // The DocumentationLanguageService does not do anything with document events
   }
@@ -143,11 +151,11 @@ package actor DocumentationLanguageService: LanguageService, Sendable {
     []
   }
 
-  package func symbolGraph(
-    forOnDiskContentsOf symbolDocumentUri: DocumentURI,
-    at location: SymbolLocation
+  package func symbolGraphForDocumentOnDisk(
+    at location: SymbolLocation,
+    manager: OnDiskDocumentManager
   ) async throws -> String? {
-    return nil
+    nil
   }
 
   package func openGeneratedInterface(

Sources/SourceKitLSP/LanguageService.swift

@@ -144,6 +144,12 @@ package protocol LanguageService: AnyObject, Sendable {
   /// Sent to close a document on the Language Server.
   func closeDocument(_ notification: DidCloseTextDocumentNotification) async
 
+  /// Sent to open up a document on the Language Server whose contents are on-disk.
+  func openDocumentOnDisk(snapshot: DocumentSnapshot, originalFile: DocumentURI) async throws
+
+  /// Sent to close a document on the Language Server whose contents are on-disk.
+  func closeDocumentOnDisk(snapshot: DocumentSnapshot) async
+
   /// Re-open the given document, discarding any in-memory state and forcing an AST to be re-built after build settings
   /// have been changed. This needs to be handled via a notification to ensure that no other request for this document
   /// is executing at the same time.
@@ -183,10 +189,7 @@ package protocol LanguageService: AnyObject, Sendable {
 
   /// Return the symbol graph at the given location for the contents of the document as they are on-disk (opposed to the
   /// in-memory modified version of the document).
-  func symbolGraph(
-    forOnDiskContentsOf symbolDocumentUri: DocumentURI,
-    at location: SymbolLocation
-  ) async throws -> String?
+  func symbolGraphForDocumentOnDisk(at location: SymbolLocation, manager: OnDiskDocumentManager) async throws -> String?
 
   /// Request a generated interface of a module to display in the IDE.
   ///

Sources/SourceKitLSP/OnDiskDocumentManager.swift

@@ -0,0 +1,75 @@
+import Foundation
+package import LanguageServerProtocol
+import SKLogging
+import SKUtilities
+import SwiftExtensions
+
+package actor OnDiskDocumentManager {
+  private weak var sourceKitLSPServer: SourceKitLSPServer?
+  private var openSnapshots: [DocumentURI: DocumentSnapshot]
+
+  fileprivate init(sourceKitLSPServer: SourceKitLSPServer) {
+    self.sourceKitLSPServer = sourceKitLSPServer
+    openSnapshots = [:]
+  }
+
+  /// Retrieves the ``LanguageService`` for a given ``DocumentURI`` and ``Language``.
+  package func languageService(for uri: DocumentURI, _ language: Language) async throws -> LanguageService {
+    guard let sourceKitLSPServer,
+      let workspace = await sourceKitLSPServer.workspaceForDocument(uri: uri),
+      let languageService = await sourceKitLSPServer.languageService(for: uri, language, in: workspace)
+    else {
+      throw ResponseError.unknown("Unable to find language service for URI: \(uri)")
+    }
+    return languageService
+  }
+
+  /// Opens a dummy ``DocumentSnapshot`` with contents from disk for a given ``DocumentURI`` and ``Language``.
+  ///
+  /// The snapshot will remain cached until ``closeAllDocuments()`` is called.
+  package func open(uri: DocumentURI, language: Language) async throws -> DocumentSnapshot {
+    guard let fileURL = uri.fileURL else {
+      throw ResponseError.unknown("Cannot create snapshot with on-disk contents for non-file URI \(uri.forLogging)")
+    }
+
+    if let cachedSnapshot = openSnapshots[uri] {
+      return cachedSnapshot
+    }
+
+    let snapshot = DocumentSnapshot(
+      uri: try DocumentURI(filePath: "\(UUID().uuidString)/\(fileURL.filePath)", isDirectory: false),
+      language: language,
+      version: 0,
+      lineTable: LineTable(try String(contentsOf: fileURL, encoding: .utf8))
+    )
+    try await languageService(for: uri, language).openDocumentOnDisk(snapshot: snapshot, originalFile: uri)
+    openSnapshots[uri] = snapshot
+    return snapshot
+  }
+
+  /// Close all of the ``DocumentSnapshot``s that were opened by this ``OnDiskDocumentManager``.
+  package func closeAllDocuments() async {
+    for snapshot in openSnapshots.values {
+      await orLog("Closing snapshot from on-disk contents: \(snapshot.uri.forLogging)") {
+        try await languageService(for: snapshot.uri, snapshot.language).closeDocumentOnDisk(snapshot: snapshot)
+      }
+    }
+    openSnapshots = [:]
+  }
+}
+
+package extension SourceKitLSPServer {
+  nonisolated(nonsending) func withOnDiskDocumentManager<T>(
+    _ body: (OnDiskDocumentManager) async throws -> T
+  ) async rethrows -> T {
+    let manager = OnDiskDocumentManager(sourceKitLSPServer: self)
+    do {
+      let result = try await body(manager)
+      await manager.closeAllDocuments()
+      return result
+    } catch {
+      await manager.closeAllDocuments()
+      throw error
+    }
+  }
+}

Sources/SwiftLanguageService/CMakeLists.txt

@@ -22,7 +22,6 @@ add_library(SwiftLanguageService STATIC
   SemanticRefactorCommand.swift
   DocumentFormatting.swift
   DiagnosticReportManager.swift
-  WithSnapshotFromDiskOpenedInSourcekitd.swift
   SemanticTokens.swift
   ExpandMacroCommand.swift
   Diagnostic.swift

Sources/SwiftLanguageService/DoccDocumentation.swift

@@ -35,6 +35,7 @@ extension SwiftLanguageService {
       throw ResponseError.invalidParams("A position must be provided for Swift files")
     }
     let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
+
     var moduleName: String? = nil
     var catalogURL: URL? = nil
     if let target = await workspace.buildServerManager.canonicalTarget(for: req.textDocument.uri) {
@@ -70,27 +71,19 @@ extension SwiftLanguageService {
       throw ResponseError.internalError("Unable to retrieve symbol graph for the document")
     }
     // Locate the documentation extension and include it in the request if one exists
-    let markupExtensionFile = await orLog("Finding markup extension file for symbol \(symbolUSR)") {
-      try await findMarkupExtensionFile(
-        workspace: workspace,
-        documentationManager: documentationManager,
-        catalogURL: catalogURL,
-        for: symbolUSR,
-        fetchSymbolGraph: { symbolLocation in
-          try await withSnapshotFromDiskOpenedInSourcekitd(
-            uri: symbolLocation.documentUri,
-            fallbackSettingsAfterTimeout: false
-          ) { (snapshot, compileCommand) in
-            let (_, _, symbolGraph) = try await self.cursorInfo(
-              snapshot,
-              compileCommand: compileCommand,
-              Range(snapshot.position(of: symbolLocation)),
-              includeSymbolGraph: true
-            )
-            return symbolGraph
+    let markupExtensionFile = await sourceKitLSPServer.withOnDiskDocumentManager { onDiskDocumentManager in
+      await orLog("Finding markup extension file for symbol \(symbolUSR)") {
+        try await findMarkupExtensionFile(
+          workspace: workspace,
+          documentationManager: documentationManager,
+          catalogURL: catalogURL,
+          for: symbolUSR,
+          fetchSymbolGraph: {
+            try await onDiskDocumentManager.languageService(for: $0.documentUri, .swift)
+              .symbolGraphForDocumentOnDisk(at: $0, manager: onDiskDocumentManager)
           }
-        }
-      )
+        )
+      }
     }
     return try await documentationManager.renderDocCDocumentation(
       symbolUSR: symbolUSR,
@@ -114,11 +107,12 @@ extension SwiftLanguageService {
     }
     let catalogIndex = try await documentationManager.catalogIndex(for: catalogURL)
     guard let index = workspace.index(checkedFor: .deletedFiles),
-      let symbolInformation = try await index.doccSymbolInformation(
-        ofUSR: symbolUSR,
-        fetchSymbolGraph: fetchSymbolGraph
-      ),
-      let markupExtensionFileURL = catalogIndex.documentationExtension(for: symbolInformation)
+      let markupExtensionFileURL = try await catalogIndex.documentationExtension(
+        for: index.doccSymbolInformation(
+          ofUSR: symbolUSR,
+          fetchSymbolGraph: fetchSymbolGraph
+        )
+      )
     else {
       return nil
     }

Sources/SwiftLanguageService/SwiftLanguageService.swift

@@ -586,6 +586,34 @@ extension SwiftLanguageService {
     }
   }
 
+  package func openDocumentOnDisk(snapshot: DocumentSnapshot, originalFile: DocumentURI) async throws {
+    let patchedCompileCommand: SwiftCompileCommand? =
+      if let buildSettings = await self.buildSettings(
+        for: originalFile,
+        fallbackAfterTimeout: false
+      ) {
+        SwiftCompileCommand(buildSettings.patching(newFile: snapshot.uri, originalFile: originalFile))
+      } else {
+        nil
+      }
+
+    _ = try await send(
+      sourcekitdRequest: \.editorOpen,
+      self.openDocumentSourcekitdRequest(snapshot: snapshot, compileCommand: patchedCompileCommand),
+      snapshot: snapshot
+    )
+  }
+
+  package func closeDocumentOnDisk(snapshot: DocumentSnapshot) async {
+    await orLog("Close on-disk document in sourcekitd '\(snapshot.uri)'") {
+      _ = try await send(
+        sourcekitdRequest: \.editorClose,
+        self.closeDocumentSourcekitdRequest(uri: snapshot.uri),
+        snapshot: snapshot
+      )
+    }
+  }
+
   /// Cancels any in-flight tasks to send a `PublishedDiagnosticsNotification` after edits.
   private func cancelInFlightPublishDiagnosticsTask(for document: DocumentURI) {
     if let inFlightTask = inFlightPublishDiagnosticsTasks[document] {