Merge pull request #2226 from matthewbastien/docc-cache-snapshots

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

Author: ahoppen

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -795,6 +795,28 @@ package actor BuildServerManager: QueueBasedMessageHandler {
     return languageFromBuildServer ?? Language(inferredFromFileExtension: document)
   }
 
+  /// Returns the language that a document should be interpreted in for background tasks where the editor doesn't
+  /// specify the document's language.
+  ///
+  /// If the language could not be determined, this method throws an error.
+  package func defaultLanguageInCanonicalTarget(for document: DocumentURI) async throws -> Language {
+    struct UnableToInferLanguage: Error, CustomStringConvertible {
+      let document: DocumentURI
+      var description: String { "Unable to infer language for \(document)" }
+    }
+
+    guard let canonicalTarget = await self.canonicalTarget(for: document) else {
+      guard let language = Language(inferredFromFileExtension: document) else {
+        throw UnableToInferLanguage(document: document)
+      }
+      return language
+    }
+    guard let language = await defaultLanguage(for: document, in: canonicalTarget) else {
+      throw UnableToInferLanguage(document: document)
+    }
+    return language
+  }
+
   /// Retrieve information about the given source file within the build server.
   package func sourceFileInfo(for document: DocumentURI) async -> SourceFileInfo? {
     return await orLog("Getting targets for source file") {

Sources/ClangLanguageService/ClangLanguageService.swift

@@ -496,14 +496,10 @@ extension ClangLanguageService {
   }
 
   package func doccDocumentation(_ req: DoccDocumentationRequest) async throws -> DoccDocumentationResponse {
-    guard let sourceKitLSPServer else {
-      throw ResponseError.unknown("Connection to the editor closed")
+    guard let language = openDocuments[req.textDocument.uri] else {
+      throw ResponseError.requestFailed("Documentation preview is not available for clang files")
     }
-
-    let snapshot = try sourceKitLSPServer.documentManager.latestSnapshot(req.textDocument.uri)
-    throw ResponseError.requestFailed(
-      "Documentation preview is not available for \(snapshot.language.description) files"
-    )
+    throw ResponseError.requestFailed("Documentation preview is not available for \(language.description) files")
   }
 
   package func symbolInfo(_ req: SymbolInfoRequest) async throws -> [SymbolDetails] {

Sources/DocumentationLanguageService/DoccDocumentationError.swift

@@ -23,7 +23,7 @@ enum DocCDocumentationError: LocalizedError {
     case .unsupportedLanguage(let language):
       return "Documentation preview is not available for \(language.description) files"
     case .indexNotAvailable:
-      return "The index is not availble to complete the request"
+      return "The index is not available to complete the request"
     case .symbolNotFound(let symbolName):
       return "Could not find symbol \(symbolName) in the project"
     }

Sources/DocumentationLanguageService/DoccDocumentationHandler.swift

@@ -104,47 +104,35 @@ 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 = await sourceKitLSPServer.workspaceForDocument(uri: location.documentUri) else {
-              throw ResponseError.internalError("Unable to find 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: { location in
+              return try await sourceKitLSPServer.primaryLanguageService(
+                for: location.documentUri,
+                workspace.buildServerManager.defaultLanguageInCanonicalTarget(for: location.documentUri),
+                in: workspace
+              )
+              .symbolGraph(forOnDiskContentsAt: location, in: workspace, manager: onDiskDocumentManager)
             }
-            let languageService = try await sourceKitLSPServer.primaryLanguageService(
-              for: location.documentUri,
-              .swift,
-              in: symbolWorkspace
-            )
-            return try await languageService.symbolGraph(
-              forOnDiskContentsOf: location.documentUri,
-              at: location
-            )
-          }
+          )
+        else {
+          throw ResponseError.requestFailed(doccDocumentationError: .symbolNotFound(symbolName))
+        }
+        let symbolGraph = try await sourceKitLSPServer.primaryLanguageService(
+          for: symbolOccurrence.location.documentUri,
+          workspace.buildServerManager.defaultLanguageInCanonicalTarget(for: symbolOccurrence.location.documentUri),
+          in: workspace
+        ).symbolGraph(forOnDiskContentsAt: symbolOccurrence.location, in: workspace, manager: onDiskDocumentManager)
+        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 = await sourceKitLSPServer.workspaceForDocument(uri: symbolDocumentUri) else {
-        throw ResponseError.internalError("Unable to find language service for \(symbolDocumentUri)")
-      }
-      let languageService = try await sourceKitLSPServer.primaryLanguageService(
-        for: symbolDocumentUri,
-        .swift,
-        in: symbolWorkspace
-      )
-      let symbolGraph = try await languageService.symbolGraph(
-        forOnDiskContentsOf: symbolDocumentUri,
-        at: symbolOccurrence.location
-      )
-      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.
@@ -175,24 +163,26 @@ extension DocumentationLanguageService {
       in: workspace
     ).symbolGraph(for: snapshot, at: position)
     // 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: { location in
-          guard let symbolWorkspace = await sourceKitLSPServer.workspaceForDocument(uri: location.documentUri) else {
-            throw ResponseError.internalError("Unable to find language service for \(location.documentUri)")
+    let markupExtensionFile = await sourceKitLSPServer.withOnDiskDocumentManager {
+      [documentationManager, documentManager = try documentManager] onDiskDocumentManager in
+      await orLog("Finding markup extension file for symbol \(symbolUSR)") {
+        try await Self.findMarkupExtensionFile(
+          workspace: workspace,
+          documentationManager: documentationManager,
+          documentManager: documentManager,
+          catalogURL: catalogURL,
+          for: symbolUSR,
+          fetchSymbolGraph: { location in
+            try await sourceKitLSPServer.primaryLanguageService(
+              for: location.documentUri,
+              snapshot.language,
+              in: workspace
+            )
+            .symbolGraph(forOnDiskContentsAt: location, in: workspace, manager: onDiskDocumentManager)
+
           }
-          let languageService = try await sourceKitLSPServer.primaryLanguageService(
-            for: location.documentUri,
-            .swift,
-            in: symbolWorkspace
-          )
-          return try await languageService.symbolGraph(forOnDiskContentsOf: location.documentUri, at: location)
-        }
-      )
+        )
+      }
     }
     return try await documentationManager.renderDocCDocumentation(
       symbolUSR: symbolUSR,
@@ -204,9 +194,10 @@ extension DocumentationLanguageService {
     )
   }
 
-  private func findMarkupExtensionFile(
+  private static func findMarkupExtensionFile(
     workspace: Workspace,
     documentationManager: DocCDocumentationManager,
+    documentManager: DocumentManager,
     catalogURL: URL?,
     for symbolUSR: String,
     fetchSymbolGraph: @Sendable (SymbolLocation) async throws -> String?
@@ -215,16 +206,17 @@ extension DocumentationLanguageService {
       return nil
     }
     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)
-    else {
+    guard let index = workspace.index(checkedFor: .deletedFiles) else {
+      return nil
+    }
+    let symbolInformation = try await index.doccSymbolInformation(
+      ofUSR: symbolUSR,
+      fetchSymbolGraph: fetchSymbolGraph
+    )
+    guard let markupExtensionFileURL = catalogIndex.documentationExtension(for: symbolInformation) else {
       return nil
     }
-    return try? documentManager.latestSnapshotOrDisk(
+    return documentManager.latestSnapshotOrDisk(
       DocumentURI(markupExtensionFileURL),
       language: .markdown
     )?.text

Sources/DocumentationLanguageService/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 {
   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

Sources/SourceKitLSP/LanguageService.swift

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+package import BuildServerIntegration
 import Foundation
 package import IndexStoreDB
 package import LanguageServerProtocol
@@ -148,6 +149,18 @@ 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.
+  ///
+  /// The snapshot will have a synthesized name and the caller is responsible for synthesizing build settings for it.
+  ///
+  /// - Important: This should only be called by `OnDiskDocumentManager`.
+  func openOnDiskDocument(snapshot: DocumentSnapshot, buildSettings: FileBuildSettings) async throws
+
+  /// Sent to close a document that was opened by `openOnDiskDocument`.
+  ///
+  /// - Important: This should only be called by `OnDiskDocumentManager`.
+  func closeOnDiskDocument(uri: DocumentURI) async throws
+
   /// 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.
@@ -197,8 +210,9 @@ 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
+    forOnDiskContentsAt location: SymbolLocation,
+    in workspace: Workspace,
+    manager: OnDiskDocumentManager
   ) async throws -> String
 
   /// Request a generated interface of a module to display in the IDE.
@@ -330,6 +344,14 @@ package extension LanguageService {
 
   func clientInitialized(_ initialized: LanguageServerProtocol.InitializedNotification) async {}
 
+  func openOnDiskDocument(snapshot: DocumentSnapshot, buildSettings: FileBuildSettings) async throws {
+    throw ResponseError.unknown("\(#function) not implemented in \(Self.self)")
+  }
+
+  func closeOnDiskDocument(uri: DocumentURI) async throws {
+    throw ResponseError.unknown("\(#function) not implemented in \(Self.self)")
+  }
+
   func willSaveDocument(_ notification: WillSaveTextDocumentNotification) async {}
 
   func didSaveDocument(_ notification: DidSaveTextDocumentNotification) async {}
@@ -368,8 +390,9 @@ package extension LanguageService {
   }
 
   func symbolGraph(
-    forOnDiskContentsOf symbolDocumentUri: DocumentURI,
-    at location: SymbolLocation
+    forOnDiskContentsAt location: SymbolLocation,
+    in workspace: Workspace,
+    manager: OnDiskDocumentManager
   ) async throws -> String {
     throw ResponseError.internalError("\(#function) not implemented in \(Self.self)")
   }