Merge pull request #2246 from ahoppen/multiple-lang-services-for-doc

Allow multiple language services to handle a single document

Author: ahoppen

Sources/DocumentationLanguageService/DoccDocumentationHandler.swift

@@ -63,32 +63,29 @@ extension DocumentationLanguageService {
           let symbolOccurrence = try await index.primaryDefinitionOrDeclarationOccurrence(
             ofDocCSymbolLink: symbolLink,
             fetchSymbolGraph: { location in
-              guard let symbolWorkspace = try await workspaceForDocument(uri: location.documentUri),
-                let languageService = await sourceKitLSPServer.languageService(
-                  for: location.documentUri,
-                  .swift,
-                  in: symbolWorkspace
-                )
-              else {
+              guard let symbolWorkspace = try await workspaceForDocument(uri: location.documentUri) else {
                 throw ResponseError.internalError("Unable to find language service for \(location.documentUri)")
               }
+              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 symbolDocumentUri = symbolOccurrence.location.documentUri
-        guard
-          let symbolWorkspace = try await workspaceForDocument(uri: symbolDocumentUri),
-          let languageService = await sourceKitLSPServer.languageService(
-            for: symbolDocumentUri,
-            .swift,
-            in: symbolWorkspace
-          )
-        else {
+        guard let symbolWorkspace = try await 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

Sources/DocumentationLanguageService/DocumentationLanguageService.swift

@@ -51,17 +51,6 @@ package actor DocumentationLanguageService: LanguageService, Sendable {
     return await sourceKitLSPServer.workspaceForDocument(uri: uri)
   }
 
-  func languageService(
-    for uri: DocumentURI,
-    _ language: LanguageServerProtocol.Language,
-    in workspace: Workspace
-  ) async throws -> LanguageService? {
-    guard let sourceKitLSPServer else {
-      throw ResponseError.unknown("Connection to the editor closed")
-    }
-    return await sourceKitLSPServer.languageService(for: uri, language, in: workspace)
-  }
-
   package nonisolated func canHandle(workspace: Workspace, toolchain: Toolchain) -> Bool {
     return true
   }

Sources/LanguageServerProtocol/Error.swift

@@ -85,6 +85,8 @@ public struct ErrorCode: RawRepresentable, Codable, Hashable, Sendable {
   public static let workspaceNotOpen: ErrorCode = ErrorCode(rawValue: -32003)
 
   /// The method is not implemented in this `LanguageService`.
+  ///
+  /// This informs `SourceKitLSPServer` that it should query secondary language services for the results.
   public static let requestNotImplemented: ErrorCode = ErrorCode(rawValue: -32004)
 }
 

Sources/SourceKitLSP/LanguageServiceRegistry.swift

@@ -34,27 +34,35 @@ struct LanguageServiceType: Hashable {
 /// Registry in which conformers to `LanguageService` can be registered to server semantic functionality for a set of
 /// languages.
 package struct LanguageServiceRegistry {
-  private var byLanguage: [Language: LanguageService.Type] = [:]
+  private var byLanguage: [Language: [LanguageServiceType]] = [:]
 
   package init() {}
 
   package mutating func register(_ languageService: LanguageService.Type, for languages: [Language]) {
     for language in languages {
-      if let existingLanguageService = byLanguage[language] {
-        logger.fault(
-          "Cannot register \(languageService) for \(language, privacy: .public) because \(existingLanguageService) is already registered"
-        )
+      let services = byLanguage[language] ?? []
+      if services.contains(LanguageServiceType(languageService)) {
+        logger.fault("\(languageService) already registered for \(language, privacy: .public)")
         continue
       }
-      byLanguage[language] = languageService
+      byLanguage[language, default: []].append(LanguageServiceType(languageService))
     }
   }
 
-  func languageService(for language: Language) -> LanguageService.Type? {
-    return byLanguage[language]
+  /// The language services that can handle a document of the given language.
+  ///
+  /// Multiple language services may be able to handle a document. Depending on the use case, callers need to combine
+  /// the results of the language services.
+  /// If it is possible to merge the results of the language service (eg. combining code actions from multiple language
+  /// services), that's the preferred choice.
+  /// Otherwise the language services occurring early in the array should be given precedence and the results of the
+  /// first language service that produces some should be returned.
+  func languageServices(for language: Language) -> [LanguageService.Type] {
+    return byLanguage[language]?.map(\.type) ?? []
   }
 
+  /// All language services that are registered in the registry.
   var languageServices: Set<LanguageServiceType> {
-    return Set(byLanguage.values.map { LanguageServiceType($0) })
+    return Set(byLanguage.values.flatMap { $0 })
   }
 }

Sources/SourceKitLSP/Rename.swift

@@ -124,7 +124,9 @@ extension SourceKitLSPServer {
     guard let snapshot = self.documentManager.latestSnapshotOrDisk(uri, language: .swift) else {
       return nil
     }
-    let swiftLanguageService = await self.languageService(for: uri, .swift, in: workspace) as? NameTranslatorService
+    let swiftLanguageService = await orLog("Getting NameTranslatorService") {
+      try await self.primaryLanguageService(for: uri, .swift, in: workspace) as? NameTranslatorService
+    }
     guard let swiftLanguageService else {
       return nil
     }
@@ -210,7 +212,7 @@ extension SourceKitLSPServer {
       return CrossLanguageName(clangName: definitionName, swiftName: swiftName, definitionLanguage: definitionLanguage)
     case .swift:
       guard
-        let swiftLanguageService = await self.languageService(
+        let swiftLanguageService = try await self.primaryLanguageService(
           for: definitionDocumentUri,
           definitionLanguage,
           in: workspace
@@ -278,9 +280,7 @@ extension SourceKitLSPServer {
     guard let workspace = await workspaceForDocument(uri: uri) else {
       throw ResponseError.workspaceNotOpen(uri)
     }
-    guard let primaryFileLanguageService = workspace.documentService(for: uri) else {
-      return nil
-    }
+    let primaryFileLanguageService = try await primaryLanguageService(for: uri, snapshot.language, in: workspace)
 
     // Determine the local edits and the USR to rename
     let renameResult = try await primaryFileLanguageService.rename(request)
@@ -395,7 +395,10 @@ extension SourceKitLSPServer {
           logger.error("Failed to get document snapshot for \(uri.forLogging)")
           return nil
         }
-        guard let languageService = await self.languageService(for: uri, language, in: workspace) else {
+        let languageService = await orLog("Getting language service to compute edits in file") {
+          try await self.primaryLanguageService(for: uri, language, in: workspace)
+        }
+        guard let languageService else {
           return nil
         }