Allow multiple files to be indexed in a single invocation using the -index-file driver mode

This allows us to index mulitple files in a single frontend invocation during background indexing, allowing us to share common work like module loading.

Author: ahoppen

Sources/SwiftDriver/Jobs/CompileJob.swift

@@ -107,7 +107,7 @@ extension Driver {
   /// corresponding primary set of outputs and, if not identical, the output
   /// paths to record in the index data (empty otherwise).
   mutating func addCompileInputs(primaryInputs: [TypedVirtualPath],
-                                 indexFilePath: TypedVirtualPath?,
+                                 indexFilePaths: [TypedVirtualPath],
                                  inputs: inout [TypedVirtualPath],
                                  inputOutputMap: inout [TypedVirtualPath: [TypedVirtualPath]],
                                  outputType: FileType?,
@@ -149,11 +149,11 @@ extension Driver {
       assert(!primaryInputs.isEmpty)
       usesPrimaryFileInputs = true
       primaryInputFiles = primaryInputs
-    } else if let path = indexFilePath {
+    } else if !indexFilePaths.isEmpty {
       // If -index-file is used, we perform a single compile but pass the
       // -index-file-path as a primary input file.
       usesPrimaryFileInputs = true
-      primaryInputFiles = [path]
+      primaryInputFiles = indexFilePaths
     } else {
       usesPrimaryFileInputs = false
       primaryInputFiles = []
@@ -243,17 +243,14 @@ extension Driver {
     commandLine.appendFlag("-frontend")
     addCompileModeOption(outputType: outputType, commandLine: &commandLine)
 
-    let indexFilePath: TypedVirtualPath?
-    if let indexFileArg = parsedOptions.getLastArgument(.indexFilePath)?.asSingle {
-      let path = try VirtualPath(path: indexFileArg)
-      indexFilePath = inputFiles.first { $0.file == path }
-    } else {
-      indexFilePath = nil
-    }
+    let indexFileArgs = try Set(
+      parsedOptions.arguments(for: .indexFilePath).map { try VirtualPath(path: $0.argument.asSingle) }
+    )
+    let indexFilePaths = inputFiles.filter { indexFileArgs.contains($0.file) }
 
     let (primaryOutputs, primaryIndexUnitOutputs) =
       try addCompileInputs(primaryInputs: primaryInputs,
-                           indexFilePath: indexFilePath,
+                           indexFilePaths: indexFilePaths,
                            inputs: &inputs,
                            inputOutputMap: &inputOutputMap,
                            outputType: outputType,
@@ -281,7 +278,7 @@ extension Driver {
       moduleOutputInfo: self.moduleOutputInfo,
       moduleOutputPaths: self.moduleOutputPaths,
       includeModuleTracePath: emitModuleTrace,
-      indexFilePath: indexFilePath)
+      indexFilePaths: indexFilePaths)
 
     // Forward migrator flags.
     try commandLine.appendLast(.apiDiffDataFile, from: &parsedOptions)

Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift

@@ -665,7 +665,7 @@ extension Driver {
                                                         moduleOutputInfo: ModuleOutputInfo,
                                                         moduleOutputPaths: SupplementalModuleTargetOutputPaths,
                                                         includeModuleTracePath: Bool,
-                                                        indexFilePath: TypedVirtualPath?) throws -> [TypedVirtualPath] {
+                                                        indexFilePaths: [TypedVirtualPath]) throws -> [TypedVirtualPath] {
     var flaggedInputOutputPairs: [(flag: String, input: TypedVirtualPath?, output: TypedVirtualPath)] = []
 
     /// Add output of a particular type, if needed.
@@ -879,7 +879,10 @@ extension Driver {
       }
       // To match the legacy driver behavior, make sure we add an entry for the
       // file under indexing and the primary output file path.
-      if let indexFilePath = indexFilePath, let idxOutput = inputOutputMap[indexFilePath]?.first {
+      for indexFilePath in indexFilePaths {
+        guard let idxOutput = inputOutputMap[indexFilePath]?.first else {
+          continue
+        }
         entries[indexFilePath.fileHandle] = [.indexData: idxOutput.fileHandle]
       }
       let outputFileMap = OutputFileMap(entries: entries)

Tests/SwiftDriverTests/SwiftDriverTests.swift

@@ -484,6 +484,40 @@ final class SwiftDriverTests: XCTestCase {
     }
   }
 
+    func testIndexMultipleFilesInSingleCommandLineInvocation() throws {
+      try withTemporaryFile { outputFileMap in
+        let outputMapContents = ByteString("""
+        {
+          "first.swift": {
+            "index-unit-output-path": "first.o"
+          },
+          "second.swift": {
+            "index-unit-output-path": "second.o"
+          }
+        }
+        """.utf8)
+        try localFileSystem.writeFileContents(outputFileMap.path, bytes: outputMapContents)
+        try assertNoDriverDiagnostics(args:
+          "swiftc", "-index-file",
+          "first.swift", "second.swift", "third.swift",
+          "-index-file-path", "first.swift",
+          "-index-file-path", "second.swift",
+          "-index-store-path", "/tmp/idx",
+          "-output-file-map", outputFileMap.path.pathString
+        ) { driver in
+            let jobs = try driver.planBuild()
+            XCTAssertEqual(jobs.count, 1)
+            let commandLine = jobs[0].commandLine
+            XCTAssertJobInvocationMatches(jobs[0], .flag("-index-unit-output-path"), .path(.relative(try RelativePath(validating: "first.o"))))
+            XCTAssertJobInvocationMatches(jobs[0], .flag("-index-unit-output-path"), .path(.relative(try RelativePath(validating: "second.o"))))
+            XCTAssertEqual(commandLine.filter { $0 == .flag("-index-unit-output-path") }.count, 2)
+            XCTAssertJobInvocationMatches(jobs[0], .flag("-primary-file"), .path(.relative(try RelativePath(validating: "first.swift"))))
+            XCTAssertJobInvocationMatches(jobs[0], .flag("-primary-file"), .path(.relative(try RelativePath(validating: "second.swift"))))
+            XCTAssertEqual(commandLine.filter { $0 == .flag("-primary-file") }.count, 2)
+        }
+      }
+  }
+
   func testMultiThreadingOutputs() throws {
     try assertDriverDiagnostics(args: "swiftc", "-c", "foo.swift", "bar.swift", "-o", "bar.ll", "-o", "foo.ll", "-num-threads", "2", "-whole-module-optimization") {
       $1.expect(.error("cannot specify -o when generating multiple output files"))
@@ -4709,7 +4743,7 @@ final class SwiftDriverTests: XCTestCase {
     // in addition to the usual flag.
     try withTemporaryDirectory { path in
       let completePath: AbsolutePath = path.appending(component: "profile.profdata")
-      
+
       try localFileSystem.writeFileContents(completePath, bytes: .init())
       var driver = try Driver(args: ["swiftc", "foo.swift",
         "-working-directory", path.pathString,
@@ -7516,7 +7550,7 @@ final class SwiftDriverTests: XCTestCase {
       XCTAssertEqual(try getLibraryLevel(flags), .spi)
     }
     try withTemporaryFile { file in
-      try localFileSystem.writeFileContents(file.path, bytes: 
+      try localFileSystem.writeFileContents(file.path, bytes:
         "// swift-module-flags: -target arm64e-apple-macos12.0"
       )
       let flags = try getAllModuleFlags(VirtualPath.absolute(file.path))
@@ -8332,7 +8366,7 @@ final class SwiftDriverTests: XCTestCase {
     XCTAssertTrue(plannedJobs[0].commandLine.contains(.flag("-load-pass-plugin=/path/to/plugin")))
 #endif
   }
-    
+
   func testSupplementaryOutputFileMapUsage() throws {
     // Ensure filenames are escaped properly when using a supplementary output file map
     try withTemporaryDirectory { path in
@@ -8363,7 +8397,7 @@ final class SwiftDriverTests: XCTestCase {
         struct D {}
         """
       )
-      
+
       let sdkArgumentsForTesting = (try? Driver.sdkArgumentsForTesting()) ?? []
       let invocationArguments = ["swiftc",
                                  "-parse-as-library",