Add driver support to specify an overriding output path to record in the index data

The frontend supports this via new options -index-unit-output-path and
-index-unit-output-path-filelist that mirror -o and -output-filelist. These are
intended to allow sharing index data across builds in separate directories (so
different -o values) that are otherwise equivalent as far as the index data is
concerned (e.g. an ASAN build and a non-ASAN build) by supplying the same
-index-unit-output-path for both.

This change updates the driver to add these new options to the frontend
invocation 1) when a new "index-unit-output-path" entry is specified for one
or more input files in the -output-file-map json or 2) if -index-file is
specified, when a new -index-unit-output-path driver option is passed.

Resolves rdar://problem/74816412

Author: Nathan Hawes

Sources/SwiftDriver/Jobs/CompileJob.swift

@@ -26,6 +26,21 @@ extension Driver {
     }
   }
 
+  mutating func computeIndexUnitOutput(for input: TypedVirtualPath, outputType: FileType) -> TypedVirtualPath? {
+    guard outputType == .indexData || outputType == .object else { return nil }
+
+    if let path = outputFileMap?.existingOutput(inputFile: input.file, outputType: .indexUnitOutputPath) {
+      return TypedVirtualPath(file: path, type: outputType)
+    }
+    if outputType == .indexData {
+      if let baseOutput = parsedOptions.getLastArgument(.indexUnitOutputPath)?.asSingle,
+         let baseOutputPath = try? VirtualPath(path: baseOutput) {
+        return TypedVirtualPath(file: baseOutputPath, type: outputType)
+      }
+    }
+    return nil
+  }
+
   mutating func computePrimaryOutput(for input: TypedVirtualPath, outputType: FileType,
                                         isTopLevel: Bool) -> TypedVirtualPath {
     if let path = outputFileMap?.existingOutput(inputFile: input.file, outputType: outputType) {
@@ -79,7 +94,8 @@ extension Driver {
     case .swift, .image, .dSYM, .dependencies, .autolink, .swiftDocumentation, .swiftInterface,
          .privateSwiftInterface, .swiftSourceInfoFile, .diagnostics, .objcHeader, .swiftDeps,
          .remap, .tbd, .moduleTrace, .yamlOptimizationRecord, .bitstreamOptimizationRecord, .pcm,
-         .pch, .clangModuleMap, .jsonCompilerFeatures, .jsonTargetInfo, .jsonSwiftArtifacts, nil:
+         .pch, .clangModuleMap, .jsonCompilerFeatures, .jsonTargetInfo, .jsonSwiftArtifacts,
+         .indexUnitOutputPath, nil:
       return false
     }
   }
@@ -91,7 +107,8 @@ extension Driver {
                                  inputs: inout [TypedVirtualPath],
                                  inputOutputMap: inout [TypedVirtualPath: TypedVirtualPath],
                                  outputType: FileType?,
-                                 commandLine: inout [Job.ArgTemplate]) -> [TypedVirtualPath] {
+                                 commandLine: inout [Job.ArgTemplate])
+  -> ([TypedVirtualPath], [TypedVirtualPath]) {
     // Collect the set of input files that are part of the Swift compilation.
     let swiftInputFiles: [TypedVirtualPath] = inputFiles.filter { $0.type.isPartOfSwiftCompilation }
 
@@ -136,6 +153,8 @@ extension Driver {
 
     // Add each of the input files.
     var primaryOutputs: [TypedVirtualPath] = []
+    var primaryIndexUnitOutputs: [TypedVirtualPath] = []
+    var indexUnitOutputDiffers = false
     for input in swiftInputFiles {
       inputs.append(input)
 
@@ -160,6 +179,13 @@ extension Driver {
                                           isTopLevel: isTopLevel)
         primaryOutputs.append(output)
         inputOutputMap[input] = output
+
+        if let indexUnitOut = computeIndexUnitOutput(for: input, outputType: outputType) {
+          indexUnitOutputDiffers = true
+          primaryIndexUnitOutputs.append(indexUnitOut)
+        } else {
+          primaryIndexUnitOutputs.append(output)
+        }
       }
     }
 
@@ -172,9 +198,22 @@ extension Driver {
                                         isTopLevel: isTopLevel)
       primaryOutputs.append(output)
       inputOutputMap[input] = output
+
+      if let indexUnitOut = computeIndexUnitOutput(for: input, outputType: outputType) {
+        indexUnitOutputDiffers = true
+        primaryIndexUnitOutputs.append(indexUnitOut)
+      } else {
+        primaryIndexUnitOutputs.append(output)
+      }
+    }
+
+    if !indexUnitOutputDiffers {
+      primaryIndexUnitOutputs.removeAll()
+    } else {
+      assert(primaryOutputs.count == primaryIndexUnitOutputs.count)
     }
 
-    return primaryOutputs
+    return (primaryOutputs, primaryIndexUnitOutputs)
   }
 
   /// Form a compile job, which executes the Swift frontend to produce various outputs.
@@ -200,12 +239,13 @@ extension Driver {
       indexFilePath = nil
     }
 
-    let primaryOutputs = addCompileInputs(primaryInputs: primaryInputs,
-                                          indexFilePath: indexFilePath,
-                                          inputs: &inputs,
-                                          inputOutputMap: &inputOutputMap,
-                                          outputType: outputType,
-                                          commandLine: &commandLine)
+    let (primaryOutputs, primaryIndexUnitOutputs) =
+      addCompileInputs(primaryInputs: primaryInputs,
+                       indexFilePath: indexFilePath,
+                       inputs: &inputs,
+                       inputOutputMap: &inputOutputMap,
+                       outputType: outputType,
+                       commandLine: &commandLine)
     outputs += primaryOutputs
 
     // FIXME: optimization record arguments are added before supplementary outputs
@@ -265,6 +305,20 @@ extension Driver {
       }
     }
 
+    // Add index unit output paths if needed.
+    if !primaryIndexUnitOutputs.isEmpty {
+      if primaryOutputs.count > fileListThreshold {
+        commandLine.appendFlag(.indexUnitOutputPathFilelist)
+        let path = RelativePath(createTemporaryFileName(prefix: "index-outputs"))
+        commandLine.appendPath(.fileList(path, .list(primaryIndexUnitOutputs.map { $0.file })))
+      } else {
+        for primaryIndexUnitOutput in primaryIndexUnitOutputs {
+          commandLine.appendFlag(.indexUnitOutputPath)
+          commandLine.appendPath(primaryIndexUnitOutput.file)
+        }
+      }
+    }
+
     try commandLine.appendLast(.embedBitcodeMarker, from: &parsedOptions)
 
     // For `-index-file` mode add `-disable-typo-correction`, since the errors
@@ -380,7 +434,8 @@ extension FileType {
     case .swift, .dSYM, .autolink, .dependencies, .swiftDocumentation, .pcm,
          .diagnostics, .objcHeader, .image, .swiftDeps, .moduleTrace, .tbd,
          .yamlOptimizationRecord, .bitstreamOptimizationRecord, .swiftInterface,
-         .privateSwiftInterface, .swiftSourceInfoFile, .clangModuleMap, .jsonSwiftArtifacts:
+         .privateSwiftInterface, .swiftSourceInfoFile, .clangModuleMap, .jsonSwiftArtifacts,
+         .indexUnitOutputPath:
       fatalError("Output type can never be a primary output")
     }
   }

Sources/SwiftDriver/Utilities/FileType.swift

@@ -114,6 +114,11 @@ public enum FileType: String, Hashable, CaseIterable, Codable {
   /// The extension isn't real.
   case indexData
 
+  /// Output path to record in the indexing data store
+  ///
+  /// This is only needed for use as a key in the output file map.
+  case indexUnitOutputPath
+
   /// Optimization record.
   case yamlOptimizationRecord = "opt.yaml"
 
@@ -185,6 +190,9 @@ extension FileType: CustomStringConvertible {
     case .indexData:
       return "index-data"
 
+    case .indexUnitOutputPath:
+      return "index-unit-output-path"
+
     case .yamlOptimizationRecord:
       return "yaml-opt-record"
 
@@ -209,7 +217,8 @@ extension FileType {
          .swiftDocumentation, .pcm, .diagnostics, .objcHeader, .image,
          .swiftDeps, .moduleTrace, .tbd, .yamlOptimizationRecord, .bitstreamOptimizationRecord,
          .swiftInterface, .privateSwiftInterface, .swiftSourceInfoFile, .jsonDependencies,
-         .clangModuleMap, .jsonTargetInfo, .jsonCompilerFeatures, .jsonSwiftArtifacts:
+         .clangModuleMap, .jsonTargetInfo, .jsonCompilerFeatures, .jsonSwiftArtifacts,
+         .indexUnitOutputPath:
       return false
     }
   }
@@ -300,6 +309,8 @@ extension FileType {
       return "bitstream-opt-record"
     case .diagnostics:
       return "diagnostics"
+    case .indexUnitOutputPath:
+      return "index-unit-output-path"
     }
   }
 }
@@ -315,7 +326,8 @@ extension FileType {
       return true
     case .image, .object, .dSYM, .pch, .sib, .raw_sib, .swiftModule,
          .swiftDocumentation, .swiftSourceInfoFile, .llvmBitcode, .diagnostics,
-         .pcm, .swiftDeps, .remap, .indexData, .bitstreamOptimizationRecord:
+         .pcm, .swiftDeps, .remap, .indexData, .bitstreamOptimizationRecord,
+         .indexUnitOutputPath:
       return false
     }
   }
@@ -331,7 +343,7 @@ extension FileType {
          .swiftSourceInfoFile, .raw_sil, .raw_sib, .diagnostics, .objcHeader, .swiftDeps, .remap,
          .importedModules, .tbd, .moduleTrace, .indexData, .yamlOptimizationRecord,
          .bitstreamOptimizationRecord, .pcm, .pch, .jsonDependencies, .clangModuleMap,
-         .jsonCompilerFeatures, .jsonTargetInfo, .jsonSwiftArtifacts:
+         .jsonCompilerFeatures, .jsonTargetInfo, .jsonSwiftArtifacts, .indexUnitOutputPath:
       return false
     }
   }

Sources/SwiftOptions/Options.swift

@@ -308,6 +308,8 @@ extension Option {
   public static let indexIgnoreSystemModules: Option = Option("-index-ignore-system-modules", .flag, attributes: [.noInteractive], helpText: "Avoid indexing system modules")
   public static let indexStorePath: Option = Option("-index-store-path", .separate, attributes: [.frontend, .argumentIsPath], metaVar: "<path>", helpText: "Store indexing data to <path>")
   public static let indexSystemModules: Option = Option("-index-system-modules", .flag, attributes: [.helpHidden, .frontend, .noDriver], helpText: "Emit index data for imported serialized swift system modules")
+  public static let indexUnitOutputPath: Option = Option("-index-unit-output-path", .separate, attributes: [.frontend, .noInteractive, .argumentIsPath], metaVar: "<file>", helpText: "Record <file> as the output file in the produced index data")
+  public static let indexUnitOutputPathFilelist: Option = Option("-index-unit-output-path-filelist", .separate, attributes: [.frontend, .noDriver], helpText: "Specify unit index output paths in a file rather than on the command line")
   public static let interpret: Option = Option("-interpret", .flag, attributes: [.helpHidden, .frontend, .noDriver], helpText: "Immediate mode", group: .modes)
   public static let I: Option = Option("-I", .joinedOrSeparate, attributes: [.frontend, .argumentIsPath], helpText: "Add directory to the import search path")
   public static let i: Option = Option("-i", .flag, group: .modes)
@@ -811,6 +813,8 @@ extension Option {
       Option.indexIgnoreSystemModules,
       Option.indexStorePath,
       Option.indexSystemModules,
+      Option.indexUnitOutputPath,
+      Option.indexUnitOutputPathFilelist,
       Option.interpret,
       Option.I,
       Option.i,

Tests/SwiftDriverTests/SwiftDriverTests.swift

@@ -614,6 +614,140 @@ final class SwiftDriverTests: XCTestCase {
     }
   }
 
+  func testOutputFileMapLoadingIndexUnitOutputPath() throws {
+    let contents = """
+    {
+      "/tmp/main.swift": {
+        "object": "/tmp/build1/main.o",
+        "index-unit-output-path": "/tmp/build2/main.o",
+      },
+      "/tmp/second.swift": {
+        "object": "/tmp/build1/second.o",
+        "index-unit-output-path": "/tmp/build2/second.o",
+      }
+    }
+    """
+
+    func getFileListElements(for filelistOpt: String, job: Job) -> [VirtualPath] {
+      let optIndex = job.commandLine.firstIndex(of: .flag(filelistOpt))!
+      let value = job.commandLine[job.commandLine.index(after: optIndex)]
+      guard case let .path(.fileList(_, valueFileList)) = value else {
+        XCTFail("Argument wasn't a filelist")
+        return []
+      }
+      guard case let .list(inputs) = valueFileList else {
+        XCTFail("FileList wasn't List")
+        return []
+      }
+      return inputs
+    }
+
+    try withTemporaryFile { file in
+      try assertNoDiagnostics { diags in
+        try localFileSystem.writeFileContents(file.path) { $0 <<< contents }
+
+        // 1. Incremental mode (single primary file)
+        // a) without filelists
+        var driver = try Driver(args: [
+          "swiftc", "-c",
+          "-output-file-map", file.path.pathString,
+          "-module-name", "test", "/tmp/second.swift", "/tmp/main.swift"
+        ])
+        var jobs = try driver.planBuild()
+        print(jobs[0].commandLine)
+        print(jobs[1].commandLine)
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-o", .path(.absolute(.init("/tmp/build1/second.o")))]))
+        XCTAssertTrue(jobs[1].commandLine.contains(subsequence: ["-o", .path(.absolute(.init("/tmp/build1/main.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-index-unit-output-path", .path(.absolute(.init("/tmp/build2/second.o")))]))
+        XCTAssertTrue(jobs[1].commandLine.contains(subsequence: ["-index-unit-output-path", .path(.absolute(.init("/tmp/build2/main.o")))]))
+
+        // b) with filelists
+        driver = try Driver(args: [
+          "swiftc", "-c", "-driver-filelist-threshold=0",
+          "-output-file-map", file.path.pathString,
+          "-module-name", "test", "/tmp/second.swift", "/tmp/main.swift"
+        ])
+        jobs = try driver.planBuild()
+        XCTAssertEqual(getFileListElements(for: "-output-filelist", job: jobs[0]),
+                       [.absolute(.init("/tmp/build1/second.o"))])
+        XCTAssertEqual(getFileListElements(for: "-index-unit-output-path-filelist", job: jobs[0]),
+                       [.absolute(.init("/tmp/build2/second.o"))])
+        XCTAssertEqual(getFileListElements(for: "-output-filelist", job: jobs[1]),
+                       [.absolute(.init("/tmp/build1/main.o"))])
+        XCTAssertEqual(getFileListElements(for: "-index-unit-output-path-filelist", job: jobs[1]),
+                       [.absolute(.init("/tmp/build2/main.o"))])
+
+
+        // 2. Batch mode (two primary files)
+        // a) without filelists
+        driver = try Driver(args: [
+          "swiftc", "-c", "-enable-batch-mode", "-driver-batch-count", "1",
+          "-output-file-map", file.path.pathString,
+          "-module-name", "test", "/tmp/second.swift", "/tmp/main.swift"
+        ])
+        jobs = try driver.planBuild()
+        print(jobs[0].commandLine)
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-o", .path(.absolute(.init("/tmp/build1/second.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-o", .path(.absolute(.init("/tmp/build1/main.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-index-unit-output-path", .path(.absolute(.init("/tmp/build2/second.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-index-unit-output-path", .path(.absolute(.init("/tmp/build2/main.o")))]))
+
+        // b) with filelists
+        driver = try Driver(args: [
+          "swiftc", "-c", "-driver-filelist-threshold=0",
+          "-enable-batch-mode", "-driver-batch-count", "1",
+          "-output-file-map", file.path.pathString,
+          "-module-name", "test", "/tmp/second.swift", "/tmp/main.swift"
+        ])
+        jobs = try driver.planBuild()
+        XCTAssertEqual(getFileListElements(for: "-output-filelist", job: jobs[0]),
+                       [.absolute(.init("/tmp/build1/second.o")), .absolute(.init("/tmp/build1/main.o"))])
+        XCTAssertEqual(getFileListElements(for: "-index-unit-output-path-filelist", job: jobs[0]),
+                       [.absolute(.init("/tmp/build2/second.o")), .absolute(.init("/tmp/build2/main.o"))])
+
+        // 3. Multi-threaded WMO
+        // a) without filelists
+        driver = try Driver(args: [
+          "swiftc", "-c", "-whole-module-optimization", "-num-threads", "2",
+          "-output-file-map", file.path.pathString,
+          "-module-name", "test", "/tmp/second.swift", "/tmp/main.swift"
+        ])
+        jobs = try driver.planBuild()
+        print(jobs[0].commandLine)
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-o", .path(.absolute(.init("/tmp/build1/second.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-index-unit-output-path", .path(.absolute(.init("/tmp/build2/second.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-o", .path(.absolute(.init("/tmp/build1/main.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-index-unit-output-path", .path(.absolute(.init("/tmp/build2/main.o")))]))
+
+        // b) with filelists
+        driver = try Driver(args: [
+          "swiftc", "-c", "-driver-filelist-threshold=0",
+          "-whole-module-optimization", "-num-threads", "2",
+          "-output-file-map", file.path.pathString,
+          "-module-name", "test", "/tmp/second.swift", "/tmp/main.swift"
+        ])
+        jobs = try driver.planBuild()
+        XCTAssertEqual(getFileListElements(for: "-output-filelist", job: jobs[0]),
+                       [.absolute(.init("/tmp/build1/second.o")), .absolute(.init("/tmp/build1/main.o"))])
+        XCTAssertEqual(getFileListElements(for: "-index-unit-output-path-filelist", job: jobs[0]),
+                       [.absolute(.init("/tmp/build2/second.o")), .absolute(.init("/tmp/build2/main.o"))])
+
+        // 4. Index-file (single primary)
+        driver = try Driver(args: [
+          "swiftc", "-c", "-enable-batch-mode", "-driver-batch-count", "1",
+          "-module-name", "test", "/tmp/second.swift", "/tmp/main.swift",
+          "-index-file", "-index-file-path", "/tmp/second.swift",
+          "-disable-batch-mode", "-o", "/tmp/build1/second.o",
+          "-index-unit-output-path", "/tmp/build2/second.o"
+        ])
+        jobs = try driver.planBuild()
+        print(jobs[0].commandLine)
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-o", .path(.absolute(.init("/tmp/build1/second.o")))]))
+        XCTAssertTrue(jobs[0].commandLine.contains(subsequence: ["-index-unit-output-path", .path(.absolute(.init("/tmp/build2/second.o")))]))
+      }
+    }
+  }
+
   func testMergeModuleEmittingDependencies() throws {
     var driver1 = try Driver(args: ["swiftc", "foo.swift", "bar.swift", "-module-name", "Foo", "-emit-dependencies", "-emit-module", "-serialize-diagnostics", "-driver-filelist-threshold=9999"])
     let plannedJobs = try driver1.planBuild().removingAutolinkExtractJobs()