A couple refactorings to necessary to support multi-file indexing

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -757,10 +757,12 @@ package actor BuildServerManager: QueueBasedMessageHandler {
 
   // MARK: Build server queries
 
-  /// Returns the toolchain that should be used to process the given document.
+  /// Returns the toolchain that should be used to process the given target.
+  ///
+  /// If `target` is `nil` or the build server does not explicitly specify a toolchain for this target, the preferred
+  /// toolchain for the given language is returned.
   package func toolchain(
-    for uri: DocumentURI,
-    in target: BuildTargetIdentifier?,
+    for target: BuildTargetIdentifier?,
     language: Language
   ) async -> Toolchain? {
     let toolchainPath = await orLog("Getting toolchain from build targets") { () -> URL? in

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -91,6 +91,14 @@ private struct InProgressIndexStore {
   var fileModificationDate: Date?
 }
 
+/// The information that's needed to index a file within a given target.
+package struct FileIndexInfo: Sendable, Hashable {
+  package let file: FileToIndex
+  package let language: Language
+  package let target: BuildTargetIdentifier
+  package let outputPath: OutputPath
+}
+
 /// Status of document indexing / target preparation in `inProgressIndexAndPreparationTasks`.
 package enum IndexTaskStatus: Comparable {
   case scheduled
@@ -500,10 +508,13 @@ package final actor SemanticIndexManager {
             continue
           }
           // If this is a source file, just index it.
+          guard let language = await buildServerManager.defaultLanguage(for: uri, in: target) else {
+            continue
+          }
           didFindTargetToIndex = true
           filesToReIndex.append(
             (
-              FileIndexInfo(file: .indexableFile(uri), target: target, outputPath: outputPath),
+              FileIndexInfo(file: .indexableFile(uri), language: language, target: target, outputPath: outputPath),
               modifiedFilesIndex.modificationDate(of: uri)
             )
           )
@@ -543,10 +554,14 @@ package final actor SemanticIndexManager {
       {
         continue
       }
+      guard let language = await buildServerManager.defaultLanguage(for: uri, in: targetAndOutputPath.key) else {
+        continue
+      }
       filesToReIndex.append(
         (
           FileIndexInfo(
             file: .headerFile(header: uri, mainFile: mainFile),
+            language: language,
             target: targetAndOutputPath.key,
             outputPath: outputPath
           ),
@@ -722,16 +737,20 @@ package final actor SemanticIndexManager {
     }
   }
 
-  /// Update the index store for the given files, assuming that their targets have already been prepared.
+  /// Update the index store for the given files, assuming that their targets has already been prepared.
   private func updateIndexStore(
-    for filesAndTargets: [FileIndexInfo],
+    for fileAndOutputPaths: [FileAndOutputPath],
+    target: BuildTargetIdentifier,
+    language: Language,
     indexFilesWithUpToDateUnit: Bool,
     preparationTaskID: UUID,
     priority: TaskPriority?
   ) async {
     let taskDescription = AnyIndexTaskDescription(
       UpdateIndexStoreTaskDescription(
-        filesToIndex: filesAndTargets,
+        filesToIndex: fileAndOutputPaths,
+        target: target,
+        language: language,
         buildServerManager: self.buildServerManager,
         index: index,
         indexStoreUpToDateTracker: indexStoreUpToDateTracker,
@@ -747,7 +766,13 @@ package final actor SemanticIndexManager {
         self.indexProgressStatusDidChange()
         return
       }
-      for fileAndTarget in filesAndTargets {
+      for fileAndOutputPath in fileAndOutputPaths {
+        let fileAndTarget = FileIndexInfo(
+          file: fileAndOutputPath.file,
+          language: language,
+          target: target,
+          outputPath: fileAndOutputPath.outputPath
+        )
         switch self.inProgressIndexTasks[fileAndTarget]?.state {
         case .updatingIndexStore(let registeredTask, _):
           if registeredTask == OpaqueQueuedIndexTask(task) {
@@ -763,7 +788,13 @@ package final actor SemanticIndexManager {
       }
       self.indexProgressStatusDidChange()
     }
-    for fileAndTarget in filesAndTargets {
+    for fileAndOutputPath in fileAndOutputPaths {
+      let fileAndTarget = FileIndexInfo(
+        file: fileAndOutputPath.file,
+        language: language,
+        target: target,
+        outputPath: fileAndOutputPath.outputPath
+      )
       switch inProgressIndexTasks[fileAndTarget]?.state {
       case .waitingForPreparation(preparationTaskID, let indexTask), .preparing(preparationTaskID, let indexTask):
         inProgressIndexTasks[fileAndTarget]?.state = .updatingIndexStore(
@@ -854,13 +885,7 @@ package final actor SemanticIndexManager {
     var newIndexTasks = 0
 
     for (fileIndexInfo, fileModificationDate) in filesToIndex {
-      guard
-        let language = await buildServerManager.defaultLanguage(
-          for: fileIndexInfo.file.mainFile,
-          in: fileIndexInfo.target
-        ),
-        UpdateIndexStoreTaskDescription.canIndex(language: language)
-      else {
+      guard UpdateIndexStoreTaskDescription.canIndex(language: fileIndexInfo.language) else {
         continue
       }
 
@@ -933,17 +958,34 @@ package final actor SemanticIndexManager {
         // And after preparation is done, index the files in the targets.
         await withTaskGroup(of: Void.self) { taskGroup in
           for target in targetsBatch {
-            // TODO: Once swiftc supports indexing of multiple files in a single invocation, increase the batch size to
-            // allow it to share AST builds between multiple files within a target.
-            // (https://github.com/swiftlang/sourcekit-lsp/issues/1268)
-            for fileBatch in filesByTarget[target]!.partition(intoBatchesOfSize: 1) {
-              taskGroup.addTask {
-                await self.updateIndexStore(
-                  for: fileBatch,
-                  indexFilesWithUpToDateUnit: indexFilesWithUpToDateUnit,
-                  preparationTaskID: preparationTaskID,
-                  priority: priority
-                )
+            var filesByLanguage: [Language: [FileIndexInfo]] = [:]
+            for fileInfo in filesByTarget[target]! {
+              filesByLanguage[fileInfo.language, default: []].append(fileInfo)
+            }
+            for (language, fileInfos) in filesByLanguage {
+              // TODO: Once swiftc supports indexing of multiple files in a single invocation, increase the batch size to
+              // allow it to share AST builds between multiple files within a target.
+              // (https://github.com/swiftlang/sourcekit-lsp/issues/1268)
+              for fileBatch in fileInfos.partition(intoBatchesOfSize: 1) {
+                taskGroup.addTask {
+                  let fileAndOutputPaths: [FileAndOutputPath] = fileBatch.compactMap {
+                    guard $0.target == target else {
+                      logger.fault(
+                        "FileIndexInfo refers to different target than should be indexed \($0.target.forLogging) vs \(target.forLogging)"
+                      )
+                      return nil
+                    }
+                    return FileAndOutputPath(file: $0.file, outputPath: $0.outputPath)
+                  }
+                  await self.updateIndexStore(
+                    for: fileAndOutputPaths,
+                    target: target,
+                    language: language,
+                    indexFilesWithUpToDateUnit: indexFilesWithUpToDateUnit,
+                    preparationTaskID: preparationTaskID,
+                    priority: priority
+                  )
+                }
               }
             }
           }