[Caching] Handle emit module job correctly for swift caching

There are two situations when emiting modules are not correctly handled
when swift caching is enabled:
* When only `-emit-module` task is requested from driver, in this case,
  we are contructing a `CompileJob` but only generating module output.
  SwiftDriver still expects all swift inputs will have a cache key even
  only the first file is generating outputs and has a cache key.
* When using `-experimental-emit-module-separately`, then swift-driver
  needs to construct a `EmitModuleJob`, which wasn't taught the concept of
  caching and not producing any cache key.

rdar://127768967

Author: cachemeifyoucan

Sources/SwiftDriver/Jobs/CompileJob.swift

@@ -397,9 +397,20 @@ extension Driver {
     } else {
       displayInputs = primaryInputs
     }
-    // Only swift input files are contributing to the cache keys.
+    // The cache key for compilation is created one per input file, and each cache key contains all the output
+    // files for that specific input file. All the module level output files are attached to the cache key for
+    // the first input file. Only the input files that produce the output will have a cache key. This behavior
+    // needs to match the cache key creation logic in swift-frontend.
     let cacheContributingInputs = inputs.enumerated().reduce(into: [(TypedVirtualPath, Int)]()) { result, input in
-      if input.element.type == .swift, displayInputs.contains(input.element) {
+      guard input.element.type == .swift else { return }
+      let singleInputKey = TypedVirtualPath(file: OutputFileMap.singleInputKey, type: .swift)
+      if inputOutputMap[singleInputKey] != nil {
+        // If singleInputKey exists, that means only the first swift file produces outputs.
+        if result.isEmpty {
+          result.append((input.element, input.offset))
+        }
+      } else if !inputOutputMap[input.element, default: []].isEmpty {
+        // Otherwise, add all the inputs that produce output.
         result.append((input.element, input.offset))
       }
     }

Sources/SwiftDriver/Jobs/EmitModuleJob.swift

@@ -119,14 +119,21 @@ extension Driver {
       commandLine.appendPath(abiPath.file)
       outputs.append(abiPath)
     }
+    let cacheContributingInputs = inputs.enumerated().reduce(into: [(TypedVirtualPath, Int)]()) { result, input in
+      // only the first swift input contributes cache key to an emit module job.
+      guard result.isEmpty, input.element.type == .swift else { return }
+      result.append((input.element, input.offset))
+    }
+    let cacheKeys = try computeOutputCacheKeyForJob(commandLine: commandLine, inputs: cacheContributingInputs)
     return Job(
       moduleName: moduleOutputInfo.name,
       kind: .emitModule,
       tool: try toolchain.resolvedTool(.swiftCompiler),
       commandLine: commandLine,
       inputs: inputs,
       primaryInputs: [],
-      outputs: outputs
+      outputs: outputs,
+      outputCacheKeys: cacheKeys
     )
   }
 

Tests/SwiftDriverTests/CachingBuildTests.swift

@@ -336,6 +336,114 @@ final class CachingBuildTests: XCTestCase {
     }
   }
 
+  func testModuleOnlyJob() throws {
+    try withTemporaryDirectory { path in
+      let main = path.appending(component: "testModuleOnlyJob.swift")
+      try localFileSystem.writeFileContents(main) {
+        $0.send("import C;import E;")
+      }
+      let other = path.appending(component: "testModuleOnlyJob2.swift")
+      try localFileSystem.writeFileContents(other) {
+        $0.send("import G;")
+      }
+      let swiftModuleInterfacesPath: AbsolutePath =
+          try testInputsPath.appending(component: "ExplicitModuleBuilds")
+                            .appending(component: "Swift")
+      let cHeadersPath: AbsolutePath =
+          try testInputsPath.appending(component: "ExplicitModuleBuilds")
+                            .appending(component: "CHeaders")
+      let casPath = path.appending(component: "cas")
+      let swiftInterfacePath: AbsolutePath = path.appending(component: "testModuleOnlyJob.swiftinterface")
+      let privateSwiftInterfacePath: AbsolutePath = path.appending(component: "testModuleOnlyJob.private.swiftinterface")
+      let modulePath: AbsolutePath = path.appending(component: "testModuleOnlyJob.swiftmodule")
+      let sdkArgumentsForTesting = (try? Driver.sdkArgumentsForTesting()) ?? []
+      var driver = try Driver(args: ["swiftc",
+                                     "-target", "x86_64-apple-macosx11.0",
+                                     "-module-name", "Test",
+                                     "-I", cHeadersPath.nativePathString(escaped: true),
+                                     "-I", swiftModuleInterfacesPath.nativePathString(escaped: true),
+                                     "-emit-module-interface-path", swiftInterfacePath.nativePathString(escaped: true),
+                                     "-emit-private-module-interface-path", privateSwiftInterfacePath.nativePathString(escaped: true),
+                                     "-explicit-module-build", "-emit-module-separately-wmo", "-disable-cmo", "-Rcache-compile-job",
+                                     "-enable-library-evolution", "-O", "-whole-module-optimization",
+                                     "-cache-compile-job", "-cas-path", casPath.nativePathString(escaped: true),
+                                     "-emit-module", "-o", modulePath.nativePathString(escaped: true),
+                                     main.nativePathString(escaped: true), other.nativePathString(escaped: true)] + sdkArgumentsForTesting,
+                              env: ProcessEnv.vars,
+                              interModuleDependencyOracle: dependencyOracle)
+      let jobs = try driver.planBuild()
+      try driver.run(jobs: jobs)
+      for job in jobs {
+          XCTAssertFalse(job.outputCacheKeys.isEmpty)
+      }
+      XCTAssertFalse(driver.diagnosticEngine.hasErrors)
+
+      let scanLibPath = try XCTUnwrap(driver.getSwiftScanLibPath())
+      try dependencyOracle.verifyOrCreateScannerInstance(fileSystem: localFileSystem,
+                                                         swiftScanLibPath: scanLibPath)
+
+      let cas = try dependencyOracle.getOrCreateCAS(pluginPath: nil, onDiskPath: casPath, pluginOptions: [])
+      if let driverCAS = driver.cas {
+        XCTAssertEqual(cas, driverCAS, "CAS should only be created once")
+      } else {
+        XCTFail("Cached compilation doesn't have a CAS")
+      }
+      try checkCASForResults(jobs: jobs, cas: cas, fs: driver.fileSystem)
+    }
+  }
+
+  func testSeparateModuleJob() throws {
+    try withTemporaryDirectory { path in
+      let main = path.appending(component: "testSeparateModuleJob.swift")
+      try localFileSystem.writeFileContents(main) {
+        $0.send("import C;import E;")
+      }
+      let swiftModuleInterfacesPath: AbsolutePath =
+          try testInputsPath.appending(component: "ExplicitModuleBuilds")
+                            .appending(component: "Swift")
+      let cHeadersPath: AbsolutePath =
+          try testInputsPath.appending(component: "ExplicitModuleBuilds")
+                            .appending(component: "CHeaders")
+      let casPath = path.appending(component: "cas")
+      let swiftInterfacePath: AbsolutePath = path.appending(component: "testSeparateModuleJob.swiftinterface")
+      let privateSwiftInterfacePath: AbsolutePath = path.appending(component: "testSeparateModuleJob.private.swiftinterface")
+      let modulePath: AbsolutePath = path.appending(component: "testSeparateModuleJob.swiftmodule")
+      let sdkArgumentsForTesting = (try? Driver.sdkArgumentsForTesting()) ?? []
+      var driver = try Driver(args: ["swiftc",
+                                     "-target", "x86_64-apple-macosx11.0",
+                                     "-module-name", "Test",
+                                     "-I", cHeadersPath.nativePathString(escaped: true),
+                                     "-I", swiftModuleInterfacesPath.nativePathString(escaped: true),
+                                     "-emit-module-path", modulePath.nativePathString(escaped: true),
+                                     "-emit-module-interface-path", swiftInterfacePath.nativePathString(escaped: true),
+                                     "-emit-private-module-interface-path", privateSwiftInterfacePath.nativePathString(escaped: true),
+                                     "-explicit-module-build", "-experimental-emit-module-separately", "-Rcache-compile-job",
+                                     "-enable-library-evolution", "-O",
+                                     "-cache-compile-job", "-cas-path", casPath.nativePathString(escaped: true),
+                                     main.nativePathString(escaped: true)] + sdkArgumentsForTesting,
+                              env: ProcessEnv.vars,
+                              interModuleDependencyOracle: dependencyOracle)
+      let jobs = try driver.planBuild()
+      for job in jobs {
+          XCTAssertFalse(job.outputCacheKeys.isEmpty)
+      }
+      try driver.run(jobs: jobs)
+      XCTAssertFalse(driver.diagnosticEngine.hasErrors)
+
+      let scanLibPath = try XCTUnwrap(driver.getSwiftScanLibPath())
+      try dependencyOracle.verifyOrCreateScannerInstance(fileSystem: localFileSystem,
+                                                         swiftScanLibPath: scanLibPath)
+
+      let cas = try dependencyOracle.getOrCreateCAS(pluginPath: nil, onDiskPath: casPath, pluginOptions: [])
+      if let driverCAS = driver.cas {
+        XCTAssertEqual(cas, driverCAS, "CAS should only be created once")
+      } else {
+        XCTFail("Cached compilation doesn't have a CAS")
+      }
+      try checkCASForResults(jobs: jobs, cas: cas, fs: driver.fileSystem)
+    }
+  }
+
   /// Test generation of explicit module build jobs for dependency modules when the driver
   /// is invoked with -explicit-module-build, -verify-emitted-module-interface and -enable-library-evolution.
   func testExplicitModuleVerifyInterfaceJobs() throws {