more robust searching for DocC catalogs

Author: matthewbastien

Sources/BuildServerProtocol/Messages/BuildTargetSourcesRequest.swift

@@ -119,6 +119,26 @@ public struct SourceItemDataKind: RawRepresentable, Codable, Hashable, Sendable
 }
 
 /// **(BSP Extension)**
+
+public enum SourceKitSourceItemKind: String, Codable {
+  /// A source file that belongs to the target
+  case source = "source"
+
+  /// A header file that is clearly associated with one target.
+  ///
+  /// For example header files in SwiftPM projects are always associated to one target and SwiftPM can provide build
+  /// settings for that header file.
+  ///
+  /// In general, build systems don't need to list all header files in the `buildTarget/sources` request: Semantic
+  /// functionality for header files is usually provided by finding a main file that includes the header file and
+  /// inferring build settings from it. Listing header files in `buildTarget/sources` allows SourceKit-LSP to provide
+  /// semantic functionality for header files if they haven't been included by any main file.
+  case header = "header"
+
+  /// A SwiftDocC documentation catalog usually ending in the ".docc" extension.
+  case doccCatalog = "doccCatalog"
+}
+
 public struct SourceKitSourceItemData: LSPAnyCodable, Codable {
   /// The language of the source file. If `nil`, the language is inferred from the file extension.
   public var language: Language?
@@ -132,7 +152,14 @@ public struct SourceKitSourceItemData: LSPAnyCodable, Codable {
   /// functionality for header files is usually provided by finding a main file that includes the header file and
   /// inferring build settings from it. Listing header files in `buildTarget/sources` allows SourceKit-LSP to provide
   /// semantic functionality for header files if they haven't been included by any main file.
-  public var isHeader: Bool?
+  public var isHeader: Bool? {
+    guard let kind else {
+      return nil
+    }
+    return kind == .header
+  }
+
+  public var kind: SourceKitSourceItemKind?
 
   /// The output path that is used during indexing for this file, ie. the `-index-unit-output-path`, if it is specified
   /// in the compiler arguments or the file that is passed as `-o`, if `-index-unit-output-path` is not specified.
@@ -144,18 +171,22 @@ public struct SourceKitSourceItemData: LSPAnyCodable, Codable {
   /// `outputPathsProvider: true` in `SourceKitInitializeBuildResponseData`.
   public var outputPath: String?
 
-  public init(language: Language? = nil, isHeader: Bool? = nil, outputPath: String? = nil) {
+  public init(language: Language? = nil, kind: SourceKitSourceItemKind? = nil, outputPath: String? = nil) {
     self.language = language
-    self.isHeader = isHeader
+    self.kind = kind
     self.outputPath = outputPath
   }
 
   public init?(fromLSPDictionary dictionary: [String: LanguageServerProtocol.LSPAny]) {
     if case .string(let language) = dictionary[CodingKeys.language.stringValue] {
       self.language = Language(rawValue: language)
     }
-    if case .bool(let isHeader) = dictionary[CodingKeys.isHeader.stringValue] {
-      self.isHeader = isHeader
+    if case .string(let rawKind) = dictionary[CodingKeys.kind.stringValue] {
+      self.kind = SourceKitSourceItemKind(rawValue: rawKind)
+    }
+    // Backwards compatibility for isHeader
+    if case .bool(let isHeader) = dictionary["isHeader"], isHeader {
+      self.kind = .header
     }
     if case .string(let outputFilePath) = dictionary[CodingKeys.outputPath.stringValue] {
       self.outputPath = outputFilePath
@@ -167,8 +198,12 @@ public struct SourceKitSourceItemData: LSPAnyCodable, Codable {
     if let language {
       result[CodingKeys.language.stringValue] = .string(language.rawValue)
     }
+    if let kind {
+      result[CodingKeys.kind.stringValue] = .string(kind.rawValue)
+    }
+    // Backwards compatibility for isHeader
     if let isHeader {
-      result[CodingKeys.isHeader.stringValue] = .bool(isHeader)
+      result["isHeader"] = .bool(isHeader)
     }
     if let outputPath {
       result[CodingKeys.outputPath.stringValue] = .string(outputPath)

Sources/BuildSystemIntegration/SwiftPMBuildSystem.swift

@@ -593,16 +593,23 @@ package actor SwiftPMBuildSystem: BuiltInBuildSystem {
           kind: .file,
           generated: false,
           dataKind: .sourceKit,
-          data: SourceKitSourceItemData(isHeader: true).encodeToLSPAny()
-        )
-      }
-      sources += (swiftPMTarget.resources + swiftPMTarget.ignored + swiftPMTarget.others).map {
-        SourceItem(
-          uri: DocumentURI($0),
-          kind: $0.isDirectory ? .directory : .file,
-          generated: false
+          data: SourceKitSourceItemData(kind: .header).encodeToLSPAny()
         )
       }
+      sources += (swiftPMTarget.resources + swiftPMTarget.ignored + swiftPMTarget.others)
+        .map { (url: URL) -> SourceItem in
+          var data: SourceKitSourceItemData? = nil
+          if url.isDirectory, url.pathExtension == "docc" {
+            data = SourceKitSourceItemData(kind: .doccCatalog)
+          }
+          return SourceItem(
+            uri: DocumentURI(url),
+            kind: url.isDirectory ? .directory : .file,
+            generated: false,
+            dataKind: data != nil ? .sourceKit : nil,
+            data: data?.encodeToLSPAny()
+          )
+        }
       result.append(SourcesItem(target: target, sources: sources))
     }
     return BuildTargetSourcesResponse(items: result)

Sources/DocCDocumentation/BuildSystemIntegrationExtensions.swift

@@ -38,21 +38,15 @@ package extension BuildSystemManager {
       return nil
     }
     let sourceFiles = (try? await sourceFiles(in: [target]).flatMap(\.sources)) ?? []
-    return sourceFiles.compactMap(\.uri.fileURL?.doccCatalogURL).first
-  }
-}
-
-package extension URL {
-  var doccCatalogURL: URL? {
-    var pathComponents = self.pathComponents
-    var result = self
-    while let lastPathComponent = pathComponents.last {
-      if lastPathComponent.hasSuffix(".docc") {
-        return result
+    let catalogURLs = sourceFiles.compactMap { sourceItem -> URL? in
+      guard sourceItem.dataKind == .sourceKit,
+        let data = SourceKitSourceItemData(fromLSPAny: sourceItem.data),
+        data.kind == .doccCatalog
+      else {
+        return nil
       }
-      pathComponents.removeLast()
-      result.deleteLastPathComponent()
-    }
-    return nil
+      return sourceItem.uri.fileURL
+    }.sorted(by: { $0.absoluteString >= $1.absoluteString })
+    return catalogURLs.first
   }
 }

Sources/DocCDocumentation/DocCDocumentationManager.swift

@@ -1,5 +1,5 @@
 import BuildServerProtocol
-import BuildSystemIntegration
+package import BuildSystemIntegration
 package import Foundation
 package import IndexStoreDB
 package import LanguageServerProtocol
@@ -11,7 +11,9 @@ package struct DocCDocumentationManager: Sendable {
   private let referenceResolutionService: DocCReferenceResolutionService
   private let catalogIndexManager: DocCCatalogIndexManager
 
-  package init() {
+  private let buildSystemManager: BuildSystemManager
+
+  package init(buildSystemManager: BuildSystemManager) {
     let symbolResolutionServer = DocumentationServer(qualityOfService: .unspecified)
     doccServer = DocCServer(
       peer: symbolResolutionServer,
@@ -20,16 +22,16 @@ package struct DocCDocumentationManager: Sendable {
     catalogIndexManager = DocCCatalogIndexManager(server: doccServer)
     referenceResolutionService = DocCReferenceResolutionService()
     symbolResolutionServer.register(service: referenceResolutionService)
+    self.buildSystemManager = buildSystemManager
   }
 
   package func filesDidChange(_ events: [FileEvent]) async {
-    let affectedCatalogURLs = events.reduce(into: Set<URL>()) { affectedCatalogURLs, event in
-      guard let catalogURL = event.uri.fileURL?.doccCatalogURL else {
-        return
+    for event in events {
+      guard let catalogURL = await buildSystemManager.doccCatalog(for: event.uri) else {
+        continue
       }
-      affectedCatalogURLs.insert(catalogURL)
+      await catalogIndexManager.invalidate(catalogURL)
     }
-    await catalogIndexManager.invalidate(catalogURLs: affectedCatalogURLs)
   }
 
   package func catalogIndex(for catalogURL: URL) async throws(DocCIndexError) -> DocCCatalogIndex {

Sources/SourceKitLSP/Documentation/DoccDocumentationHandler.swift

@@ -25,7 +25,10 @@ extension DocumentationLanguageService {
     guard let sourceKitLSPServer else {
       throw ResponseError.internalError("SourceKit-LSP is shutting down")
     }
-    let documentationManager = sourceKitLSPServer.documentationManager
+    guard let workspace = await sourceKitLSPServer.workspaceForDocument(uri: req.textDocument.uri) else {
+      throw ResponseError.workspaceNotOpen(req.textDocument.uri)
+    }
+    let documentationManager = workspace.doccDocumentationManager
     let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
     guard let workspace = await sourceKitLSPServer.workspaceForDocument(uri: req.textDocument.uri) else {
       throw ResponseError.workspaceNotOpen(req.textDocument.uri)
@@ -42,7 +45,7 @@ extension DocumentationLanguageService {
       )
     case .markdown:
       if case let .symbol(symbolName) = MarkdownTitleFinder.find(parsing: snapshot.text) {
-        if let moduleName = moduleName, symbolName == moduleName {
+        if let moduleName, symbolName == moduleName {
           // This is a page representing the module itself.
           // Create a dummy symbol graph and tell SwiftDocC to convert the module name.
           let emptySymbolGraph = String(
@@ -66,53 +69,52 @@ extension DocumentationLanguageService {
             moduleName: moduleName,
             catalogURL: catalogURL
           )
-        } else {
-          // This is a symbol extension page. Find the symbol so that we can include it in the request.
-          guard let index = workspace.index(checkedFor: .deletedFiles) else {
-            throw ResponseError.requestFailed(doccDocumentationError: .indexNotAvailable)
-          }
-          guard let symbolLink = await documentationManager.symbolLink(string: symbolName),
-            let symbolOccurrence = await documentationManager.primaryDefinitionOrDeclarationOccurrence(
-              ofDocCSymbolLink: symbolLink,
-              in: index
-            )
-          else {
-            throw ResponseError.requestFailed(doccDocumentationError: .symbolNotFound(symbolName))
-          }
-          guard
-            let symbolWorkspace = try await workspaceForDocument(uri: symbolOccurrence.location.documentUri),
-            let languageService = try await languageService(
-              for: symbolOccurrence.location.documentUri,
-              .swift,
-              in: symbolWorkspace
-            ) as? SwiftLanguageService,
-            let symbolSnapshot = try documentManager.latestSnapshotOrDisk(
-              symbolOccurrence.location.documentUri,
-              language: .swift
-            )
-          else {
-            throw ResponseError.internalError(
-              "Unable to find Swift language service for \(symbolOccurrence.location.documentUri)"
-            )
-          }
-          let position = symbolSnapshot.position(of: symbolOccurrence.location)
-          let cursorInfo = try await languageService.cursorInfo(
+        }
+        // This is a symbol extension page. Find the symbol so that we can include it in the request.
+        guard let index = workspace.index(checkedFor: .deletedFiles) else {
+          throw ResponseError.requestFailed(doccDocumentationError: .indexNotAvailable)
+        }
+        guard let symbolLink = await documentationManager.symbolLink(string: symbolName),
+          let symbolOccurrence = await documentationManager.primaryDefinitionOrDeclarationOccurrence(
+            ofDocCSymbolLink: symbolLink,
+            in: index
+          )
+        else {
+          throw ResponseError.requestFailed(doccDocumentationError: .symbolNotFound(symbolName))
+        }
+        guard
+          let symbolWorkspace = try await workspaceForDocument(uri: symbolOccurrence.location.documentUri),
+          let languageService = try await languageService(
+            for: symbolOccurrence.location.documentUri,
+            .swift,
+            in: symbolWorkspace
+          ) as? SwiftLanguageService,
+          let symbolSnapshot = try documentManager.latestSnapshotOrDisk(
             symbolOccurrence.location.documentUri,
-            position..<position,
-            includeSymbolGraph: true,
-            fallbackSettingsAfterTimeout: false
+            language: .swift
           )
-          guard let symbolGraph = cursorInfo.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
+        else {
+          throw ResponseError.internalError(
+            "Unable to find Swift language service for \(symbolOccurrence.location.documentUri)"
           )
         }
+        let position = symbolSnapshot.position(of: symbolOccurrence.location)
+        let cursorInfo = try await languageService.cursorInfo(
+          symbolOccurrence.location.documentUri,
+          position..<position,
+          includeSymbolGraph: true,
+          fallbackSettingsAfterTimeout: false
+        )
+        guard let symbolGraph = cursorInfo.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 an article that can be rendered on its own
       return try await documentationManager.renderDocCDocumentation(

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -79,10 +79,6 @@ package actor SourceKitLSPServer {
 
   package let documentManager = DocumentManager()
 
-  #if canImport(DocCDocumentation)
-  package let documentationManager = DocCDocumentationManager()
-  #endif
-
   /// The `TaskScheduler` that schedules all background indexing tasks.
   ///
   /// Shared process-wide to ensure the scheduled index operations across multiple workspaces don't exceed the maximum
@@ -1476,9 +1472,6 @@ extension SourceKitLSPServer {
     // (e.g. Package.swift doesn't have build settings but affects build
     // settings). Inform the build system about all file changes.
     await workspaces.concurrentForEach { await $0.filesDidChange(notification.changes) }
-    #if canImport(DocCDocumentation)
-    await documentationManager.filesDidChange(notification.changes)
-    #endif
   }
 
   func setBackgroundIndexingPaused(_ request: SetOptionsRequest) async throws -> VoidResponse {

Sources/SourceKitLSP/Swift/DoccDocumentation.swift

@@ -23,7 +23,10 @@ extension SwiftLanguageService {
     guard let sourceKitLSPServer else {
       throw ResponseError.internalError("SourceKit-LSP is shutting down")
     }
-    let documentationManager = sourceKitLSPServer.documentationManager
+    guard let workspace = await sourceKitLSPServer.workspaceForDocument(uri: req.textDocument.uri) else {
+      throw ResponseError.workspaceNotOpen(req.textDocument.uri)
+    }
+    let documentationManager = workspace.doccDocumentationManager
     guard let position = req.position else {
       throw ResponseError.invalidParams("A position must be provided for Swift files")
     }

Sources/SourceKitLSP/Workspace.swift

@@ -26,6 +26,10 @@ import ToolchainRegistry
 import struct TSCBasic.AbsolutePath
 import struct TSCBasic.RelativePath
 
+#if canImport(DocCDocumentation)
+package import DocCDocumentation
+#endif
+
 /// Actor that caches realpaths for `sourceFilesWithSameRealpath`.
 fileprivate actor SourceFilesWithSameRealpathInferrer {
   private let buildSystemManager: BuildSystemManager
@@ -97,6 +101,10 @@ package final class Workspace: Sendable, BuildSystemManagerDelegate {
   /// The build system manager to use for documents in this workspace.
   package let buildSystemManager: BuildSystemManager
 
+  #if canImport(DocCDocumentation)
+  package let doccDocumentationManager: DocCDocumentationManager
+  #endif
+
   private let sourceFilesWithSameRealpathInferrer: SourceFilesWithSameRealpathInferrer
 
   let options: SourceKitLSPOptions
@@ -145,6 +153,9 @@ package final class Workspace: Sendable, BuildSystemManagerDelegate {
     self.options = options
     self._uncheckedIndex = ThreadSafeBox(initialValue: uncheckedIndex)
     self.buildSystemManager = buildSystemManager
+    #if canImport(DocCDocumentation)
+    self.doccDocumentationManager = DocCDocumentationManager(buildSystemManager: buildSystemManager)
+    #endif
     self.sourceFilesWithSameRealpathInferrer = SourceFilesWithSameRealpathInferrer(
       buildSystemManager: buildSystemManager
     )
@@ -392,6 +403,9 @@ package final class Workspace: Sendable, BuildSystemManagerDelegate {
 
     // Notify all clients about the reported and inferred edits.
     await buildSystemManager.filesDidChange(events)
+    #if canImport(DocCDocumentation)
+    await doccDocumentationManager.filesDidChange(events)
+    #endif
 
     async let updateSyntacticIndex: Void = await syntacticTestIndex.filesDidChange(events)
     async let updateSemanticIndex: Void? = await semanticIndexManager?.filesDidChange(events)