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

@@ -278,6 +278,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

@@ -179,6 +179,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

@@ -104,9 +104,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

@@ -119,6 +119,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

@@ -59,10 +59,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 {
@@ -123,7 +125,7 @@ public actor SwiftPMBuildSystem {
   @_spi(Testing) public let toolsBuildParameters: BuildParameters
   @_spi(Testing) public let destinationBuildParameters: BuildParameters
   private let fileSystem: FileSystem
-  private let toolchainRegistry: ToolchainRegistry
+  private let toolchain: SKCore.Toolchain
 
   private let swiftBuildSupportsPrepareForIndexingTask = SwiftExtensions.ThreadSafeBox<Task<Bool, Never>?>(
     initialValue: nil
@@ -184,7 +186,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.experimentalFeatures = experimentalFeatures
 
     guard let packageRoot = findPackageDirectory(containing: workspacePath, fileSystem) else {
@@ -193,12 +199,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),
@@ -217,7 +223,7 @@ public actor SwiftPMBuildSystem {
       fileSystem: fileSystem,
       location: location,
       configuration: configuration,
-      customHostToolchain: toolchain
+      customHostToolchain: swiftPMToolchain
     )
 
     let buildConfiguration: PackageModel.BuildConfiguration
@@ -230,17 +236,21 @@ public actor SwiftPMBuildSystem {
 
     self.toolsBuildParameters = try BuildParameters(
       destination: .host,
-      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
     )
 
     self.destinationBuildParameters = try BuildParameters(
       destination: .target,
-      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
     )
 
@@ -411,7 +421,7 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
     #endif
     // Fix up compiler arguments that point to a `/Modules` subdirectory if the Swift version in the toolchain is less
     // than 6.0 because it places the modules one level higher up.
-    let toolchainVersion = await orLog("Getting Swift version") { try await toolchainRegistry.default?.swiftVersion }
+    let toolchainVersion = await orLog("Getting Swift version") { try await toolchain.swiftVersion }
     guard let toolchainVersion, toolchainVersion < SwiftVersion(6, 0) else {
       return compileArguments
     }
@@ -471,6 +481,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.
@@ -535,7 +549,7 @@ 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, logMessageToIndexLog: logMessageToIndexLog)
+      await orLog("Preparing") { try await prepare(singleTarget: target, logMessageToIndexLog: logMessageToIndexLog) }
     }
     let filesInPreparedTargets = targets.flatMap { self.targets[$0]?.buildTarget.sources ?? [] }
     await fileDependenciesUpdatedDebouncer.scheduleCall(Set(filesInPreparedTargets.map(DocumentURI.init)))
@@ -552,16 +566,13 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
 
     // 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)")
     var arguments = [
       swift.pathString, "build",
       "--package-path", workspacePath.pathString,

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -521,7 +521,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

@@ -465,6 +465,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

@@ -57,6 +57,10 @@ actor 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")]
   }