Let the build system determine which toolchain to use for a document

This allows us to fix a toolchain when using a `SwiftPMBuildSystem`, which is critical to ensure that a target gets prepared using the same toolchain that is used to index it and that is used for sourcekitd.

Author: ahoppen

Sources/SKCore/BuildServerBuildSystem.swift

@@ -275,6 +275,10 @@ extension BuildServerBuildSystem: BuildSystem {
     return nil
   }
 
+  public func toolchain(for uri: DocumentURI, _ language: Language) async -> SKCore.Toolchain? {
+    return nil
+  }
+
   public func configuredTargets(for document: DocumentURI) async -> [ConfiguredTarget] {
     return [ConfiguredTarget(targetID: "dummy", runDestinationID: "dummy")]
   }

Sources/SKCore/BuildSystem.swift

@@ -149,6 +149,11 @@ public protocol BuildSystem: AnyObject, Sendable {
   /// If `nil` is returned, the language based on the file's extension.
   func defaultLanguage(for document: DocumentURI) async -> Language?
 
+  /// The toolchain that should be used to open the given document.
+  ///
+  /// If `nil` is returned, then the default toolchain for the given language is used.
+  func toolchain(for uri: DocumentURI, _ language: Language) async -> Toolchain?
+
   /// Register the given file for build-system level change notifications, such
   /// as command line flag changes, dependency changes, etc.
   ///

Sources/SKCore/BuildSystemManager.swift

@@ -100,9 +100,18 @@ extension BuildSystemManager {
 
   /// Returns the toolchain that should be used to process the given document.
   public func toolchain(for uri: DocumentURI, _ language: Language) async -> Toolchain? {
-    // To support multiple toolchains within a single workspace, we need to ask the build system which toolchain to use
-    // for this document.
-    return await toolchainRegistry.defaultToolchain(for: language)
+    if let toolchain = await buildSystem?.toolchain(for: uri, language) {
+      return toolchain
+    }
+
+    switch language {
+    case .swift:
+      return await toolchainRegistry.preferredToolchain(containing: [\.sourcekitd, \.swift, \.swiftc])
+    case .c, .cpp, .objective_c, .objective_cpp:
+      return await toolchainRegistry.preferredToolchain(containing: [\.clang, \.clangd])
+    default:
+      return nil
+    }
   }
 
   /// - Note: Needed so we can set the delegate from a different isolation context.

Sources/SKCore/CompilationDatabaseBuildSystem.swift

@@ -121,6 +121,10 @@ extension CompilationDatabaseBuildSystem: BuildSystem {
     return nil
   }
 
+  public func toolchain(for uri: DocumentURI, _ language: Language) async -> SKCore.Toolchain? {
+    return nil
+  }
+
   public func configuredTargets(for document: DocumentURI) async -> [ConfiguredTarget] {
     return [ConfiguredTarget(targetID: "dummy", runDestinationID: "dummy")]
   }

Sources/SKCore/ToolchainRegistry.swift

@@ -246,26 +246,14 @@ public final actor ToolchainRegistry {
     return darwinToolchainOverride ?? ToolchainRegistry.darwinDefaultToolchainIdentifier
   }
 
-  /// The toolchain to use for a document in the given language if the build system doesn't override it.
-  func defaultToolchain(for language: Language) -> Toolchain? {
-    let supportsLang = { (toolchain: Toolchain) -> Bool in
-      // FIXME: the fact that we're looking at clangd/sourcekitd instead of the compiler indicates this method needs a parameter stating what kind of tool we're looking for.
-      switch language {
-      case .swift:
-        return toolchain.sourcekitd != nil
-      case .c, .cpp, .objective_c, .objective_cpp:
-        return toolchain.clangd != nil
-      default:
-        return false
-      }
-    }
-
-    if let toolchain = self.default, supportsLang(toolchain) {
+  /// Returns the preferred toolchain that contains all the tools at the given key paths.
+  public func preferredToolchain(containing requiredTools: [KeyPath<Toolchain, AbsolutePath?>]) -> Toolchain? {
+    if let toolchain = self.default, requiredTools.allSatisfy({ toolchain[keyPath: $0] != nil }) {
       return toolchain
     }
 
     for toolchain in toolchains {
-      if supportsLang(toolchain) {
+      if requiredTools.allSatisfy({ toolchain[keyPath: $0] != nil }) {
         return toolchain
       }
     }

Sources/SKSwiftPMWorkspace/SwiftPMBuildSystem.swift

@@ -57,10 +57,12 @@ public typealias BuildServerTarget = BuildServerProtocol.BuildTarget
 
 /// Same as `toolchainRegistry.default`.
 ///
-/// Needed to work around a compiler crash that prevents us from accessing `toolchainRegistry.default` in
+/// Needed to work around a compiler crash that prevents us from accessing `toolchainRegistry.preferredToolchain` in
 /// `SwiftPMWorkspace.init`.
-private func getDefaultToolchain(_ toolchainRegistry: ToolchainRegistry) async -> SKCore.Toolchain? {
-  return await toolchainRegistry.default
+private func preferredToolchain(_ toolchainRegistry: ToolchainRegistry) async -> SKCore.Toolchain? {
+  return await toolchainRegistry.preferredToolchain(containing: [
+    \.clang, \.clangd, \.sourcekitd, \.swift, \.swiftc,
+  ])
 }
 
 fileprivate extension BuildTriple {
@@ -116,7 +118,7 @@ public actor SwiftPMBuildSystem {
   let workspace: Workspace
   public let buildParameters: BuildParameters
   let fileSystem: FileSystem
-  private let toolchainRegistry: ToolchainRegistry
+  private let toolchain: SKCore.Toolchain
 
   var fileToTargets: [AbsolutePath: [SwiftBuildTarget]] = [:]
   var sourceDirToTargets: [AbsolutePath: [SwiftBuildTarget]] = [:]
@@ -173,7 +175,11 @@ public actor SwiftPMBuildSystem {
   ) async throws {
     self.workspacePath = workspacePath
     self.fileSystem = fileSystem
-    self.toolchainRegistry = toolchainRegistry
+    guard let toolchain = await preferredToolchain(toolchainRegistry) else {
+      throw Error.cannotDetermineHostToolchain
+    }
+
+    self.toolchain = toolchain
     self.forceResolvedVersions = forceResolvedVersions
 
     guard let packageRoot = findPackageDirectory(containing: workspacePath, fileSystem) else {
@@ -182,12 +188,12 @@ public actor SwiftPMBuildSystem {
 
     self.projectRoot = try resolveSymlinks(packageRoot)
 
-    guard let destinationToolchainBinDir = await getDefaultToolchain(toolchainRegistry)?.swiftc?.parentDirectory else {
+    guard let destinationToolchainBinDir = toolchain.swiftc?.parentDirectory else {
       throw Error.cannotDetermineHostToolchain
     }
 
     let swiftSDK = try SwiftSDK.hostSwiftSDK(AbsolutePath(destinationToolchainBinDir))
-    let toolchain = try UserToolchain(swiftSDK: swiftSDK)
+    let swiftPMToolchain = try UserToolchain(swiftSDK: swiftSDK)
 
     var location = try Workspace.Location(
       forRootPackage: AbsolutePath(packageRoot),
@@ -204,7 +210,7 @@ public actor SwiftPMBuildSystem {
       fileSystem: fileSystem,
       location: location,
       configuration: configuration,
-      customHostToolchain: toolchain
+      customHostToolchain: swiftPMToolchain
     )
 
     let buildConfiguration: PackageModel.BuildConfiguration
@@ -216,9 +222,11 @@ public actor SwiftPMBuildSystem {
     }
 
     self.buildParameters = try BuildParameters(
-      dataPath: location.scratchDirectory.appending(component: toolchain.targetTriple.platformBuildPathComponent),
+      dataPath: location.scratchDirectory.appending(
+        component: swiftPMToolchain.targetTriple.platformBuildPathComponent
+      ),
       configuration: buildConfiguration,
-      toolchain: toolchain,
+      toolchain: swiftPMToolchain,
       flags: buildSetup.flags
     )
 
@@ -407,6 +415,10 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
     return nil
   }
 
+  public func toolchain(for uri: DocumentURI, _ language: Language) async -> SKCore.Toolchain? {
+    return toolchain
+  }
+
   public func configuredTargets(for uri: DocumentURI) -> [ConfiguredTarget] {
     guard let url = uri.fileURL, let path = try? AbsolutePath(validating: url.path) else {
       // We can't determine targets for non-file URIs.
@@ -444,23 +456,20 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
     // TODO (indexing): Support preparation of multiple targets at once.
     // https://github.com/apple/sourcekit-lsp/issues/1262
     for target in targets {
-      try await prepare(singleTarget: target)
+      await orLog("Preparing") { try await self.prepare(singleTarget: target) }
     }
   }
 
   private func prepare(singleTarget target: ConfiguredTarget) async throws {
     // TODO (indexing): Add a proper 'prepare' job in SwiftPM instead of building the target.
     // https://github.com/apple/sourcekit-lsp/issues/1254
-    guard let toolchain = await toolchainRegistry.default else {
-      logger.error("Not preparing because not toolchain exists")
-      return
-    }
     guard let swift = toolchain.swift else {
       logger.error(
-        "Not preparing because toolchain at \(toolchain.identifier) does not contain a Swift compiler"
+        "Not preparing because toolchain at \(self.toolchain.identifier) does not contain a Swift compiler"
       )
       return
     }
+    logger.debug("Preparing '\(target.targetID)' using \(self.toolchain.identifier)")
     let arguments = [
       swift.pathString, "build",
       "--package-path", workspacePath.pathString,

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -850,7 +850,7 @@ public actor SourceKitLSPServer {
       return nil
     }
 
-    logger.log("Using toolchain \(toolchain.displayName) (\(toolchain.identifier)) for \(uri.forLogging)")
+    logger.log("Using toolchain \(toolchain.identifier) (\(toolchain.identifier)) for \(uri.forLogging)")
 
     return workspace.documentService.withLock { documentService in
       if let concurrentlySetService = documentService[uri] {

Tests/SKCoreTests/BuildSystemManagerTests.swift

@@ -453,6 +453,10 @@ class ManualBuildSystem: BuildSystem {
     return nil
   }
 
+  public func toolchain(for uri: DocumentURI, _ language: Language) async -> SKCore.Toolchain? {
+    return nil
+  }
+
   public func configuredTargets(for document: DocumentURI) async -> [ConfiguredTarget] {
     return [ConfiguredTarget(targetID: "dummy", runDestinationID: "dummy")]
   }

Tests/SourceKitLSPTests/BuildSystemTests.swift

@@ -51,6 +51,10 @@ final class TestBuildSystem: BuildSystem {
     return nil
   }
 
+  public func toolchain(for uri: DocumentURI, _ language: Language) async -> SKCore.Toolchain? {
+    return nil
+  }
+
   public func configuredTargets(for document: DocumentURI) async -> [ConfiguredTarget] {
     return [ConfiguredTarget(targetID: "dummy", runDestinationID: "dummy")]
   }