Teach Driver to generate swift-frontend invocations to build prebuilt module cache

We should be able to generate prebuilt module cache by reusing
the infrastructure libSwiftDriver already provides. More specifically, we could
(1) write a dummy Swift file imports every frameworks/libraries in an SDK,
(2) use module dependency scanner in the driver to generate module building commands,
and (3) use MultiJobExecutor to execute those commands topologically.

This PR is for (2).

Author: nkcsgexi

Package.swift

@@ -21,6 +21,9 @@ let package = Package(
     .executable(
       name: "swift-help",
       targets: ["swift-help"]),
+    .executable(
+      name: "swift-build-sdk-interfaces",
+      targets: ["swift-build-sdk-interfaces"]),
     .library(
       name: "SwiftDriver",
       targets: ["SwiftDriver"]),
@@ -93,6 +96,11 @@ let package = Package(
       name: "swift-help",
       dependencies: ["SwiftOptions", "ArgumentParser", "SwiftToolsSupport-auto"]),
 
+    /// The help executable.
+    .target(
+      name: "swift-build-sdk-interfaces",
+      dependencies: ["SwiftDriver", "SwiftDriverExecution"]),
+
     /// The `makeOptions` utility (for importing option definitions).
     .target(
       name: "makeOptions",

Sources/SwiftDriver/Jobs/Planning.swift

@@ -496,6 +496,91 @@ extension Driver {
     return try explicitDependencyBuildPlanner!.generateExplicitModuleDependenciesBuildJobs()
   }
 
+
+  public mutating func generatePrebuitModuleGenerationJobs(_ inputMap: [String: [TypedVirtualPath]],
+                                                           _ prebuiltModuleDir: AbsolutePath) throws -> [Job] {
+    guard let sdkPath = frontendTargetInfo.sdkPath?.path else {
+      // cannot proceed without SDK path
+      return []
+    }
+    // Run the dependency scanner and update the dependency oracle with the results
+    let dependencyGraph = try gatherModuleDependencies()
+    var jobs: [Job] = []
+
+    // Create directories for each Swift module
+    try inputMap.forEach {
+      try localFileSystem.createDirectory(prebuiltModuleDir
+        .appending(RelativePath($0.key + ".swiftmodule")))
+    }
+
+    let outputMap: [String: [TypedVirtualPath]] =
+      Dictionary.init(uniqueKeysWithValues: inputMap.map { key, value in
+      let outputPaths: [TypedVirtualPath] = value.map {
+        let path = prebuiltModuleDir.appending(RelativePath(key + ".swiftmodule"))
+          .appending(RelativePath($0.file.basenameWithoutExt + ".swiftmodule"))
+        return TypedVirtualPath(file: VirtualPath.absolute(path).intern(),
+                                type: .swiftModule)
+      }
+      return (key, outputPaths)
+    })
+
+    func getDependenciesPaths(_ module: String) throws -> [TypedVirtualPath] {
+      var results: [TypedVirtualPath] = []
+      let info = dependencyGraph.modules[.swift(module)]!
+      guard let dependencies = info.directDependencies else {
+        return results
+      }
+
+      for dep in dependencies {
+        if case let .swift(moduleName) = dep {
+          if let outputs = outputMap[moduleName] {
+            results.append(contentsOf: outputs)
+          }
+        }
+      }
+      return results
+    }
+
+    let moduleInfo = dependencyGraph.mainModule
+    if let dependencies = moduleInfo.directDependencies {
+      for dep in dependencies {
+        let moduleName = dep.moduleName
+        if let inputPaths = inputMap[moduleName] {
+          let outputPaths = outputMap[moduleName]!
+          assert(inputPaths.count == outputPaths.count)
+          assert(!inputPaths.isEmpty)
+          let dependencies = try getDependenciesPaths(moduleName)
+          for i in 0..<inputPaths.count {
+            let inputPath = inputPaths[i]
+            let outputPath = outputPaths[i]
+            var commandLine: [Job.ArgTemplate] = []
+            commandLine.appendFlag(.compileModuleFromInterface)
+            commandLine.appendFlag(.sdk)
+            commandLine.append(.path(VirtualPath.lookup(sdkPath)))
+            commandLine.appendFlag(.prebuiltModuleCachePath)
+            commandLine.appendPath(prebuiltModuleDir)
+            commandLine.appendFlag(.moduleName)
+            commandLine.appendFlag(moduleName)
+            commandLine.appendFlag(.o)
+            commandLine.appendPath(outputPath.file)
+            commandLine.appendPath(inputPath.file)
+            jobs.append(Job(
+              moduleName: moduleName,
+              kind: .compile,
+              tool: .absolute(try toolchain.getToolPath(.swiftCompiler)),
+              commandLine: commandLine,
+              inputs: dependencies,
+              primaryInputs: [],
+              outputs: [outputPath]
+            ))
+          }
+        }
+      }
+    }
+
+    return jobs
+  }
+
   private mutating func gatherModuleDependencies()
   throws -> InterModuleDependencyGraph {
     var dependencyGraph = try performDependencyScan()
@@ -551,6 +636,72 @@ extension Driver {
 
 }
 
+public struct SDKPrebuiltModuleInputsCollector {
+  let sdkPath: AbsolutePath
+  let nonFrameworkDirs = [RelativePath("usr/lib/swift"),
+                          RelativePath("System/iOSSupport/usr/lib/swift")]
+  let frameworkDirs = [RelativePath("System/Library/Frameworks"),
+                      RelativePath("System/iOSSupport/System/Library/Frameworks")]
+  public init(_ sdkPath: AbsolutePath) {
+    self.sdkPath = sdkPath
+  }
+
+  public func collectSwiftInterfaceMap() throws -> [String: [TypedVirtualPath]] {
+    var results: [String: [TypedVirtualPath]] = [:]
+
+    func updateResults(_ dir: AbsolutePath) throws {
+      if !localFileSystem.exists(dir) {
+        return
+      }
+      let moduleName = dir.basenameWithoutExt
+      if results[moduleName] == nil {
+        results[moduleName] = []
+      }
+      try localFileSystem.getDirectoryContents(dir).forEach {
+        let currentFile = try VirtualPath(path: $0)
+        if currentFile.extension == "swiftinterface" {
+          let interfacePath = TypedVirtualPath(file: currentFile.intern(),
+                                              type: .swiftInterface)
+          if !results[moduleName]!.contains(interfacePath) {
+            results[moduleName]!.append(interfacePath)
+          }
+        }
+      }
+    }
+    for dir in frameworkDirs {
+      let frameDir = AbsolutePath(sdkPath, dir)
+      if !localFileSystem.exists(frameDir) {
+        continue
+      }
+      try localFileSystem.getDirectoryContents(frameDir).forEach {
+        let frameworkPath = try VirtualPath(path: $0)
+        if frameworkPath.extension != "framework" {
+          return
+        }
+        let moduleName = frameworkPath.basenameWithoutExt
+        let swiftModulePath = frameworkPath
+          .appending(component: "Modules")
+          .appending(component: moduleName + ".swiftmodule").relativePath!
+        try updateResults(AbsolutePath(frameDir, swiftModulePath))
+      }
+    }
+    for dir in nonFrameworkDirs {
+      let swiftModuleDir = AbsolutePath(sdkPath, dir)
+      if !localFileSystem.exists(swiftModuleDir) {
+        continue
+      }
+      try localFileSystem.getDirectoryContents(swiftModuleDir).forEach {
+        let swiftModulePath = try VirtualPath(path: $0).relativePath!
+        if swiftModulePath.extension != "swiftmodule" {
+          return
+        }
+        try updateResults(AbsolutePath(swiftModuleDir, swiftModulePath))
+      }
+    }
+    return results
+  }
+}
+
 /// MARK: Planning
 extension Driver {
   /// Create a job if needed for simple requests that can be immediately

Sources/swift-build-sdk-interfaces/main.swift

@@ -0,0 +1,60 @@
+import SwiftDriverExecution
+import SwiftDriver
+import TSCLibc
+import TSCBasic
+import TSCUtility
+
+let diagnosticsEngine = DiagnosticsEngine(handlers: [Driver.stderrDiagnosticsHandler])
+
+guard let sdkPath = ProcessEnv.vars["SDKROOT"] else {
+  diagnosticsEngine.emit(.error("need to set SDKROOT"))
+  exit(1)
+}
+guard let swiftcPath = ProcessEnv.vars["SWIFT_EXEC"] else {
+  diagnosticsEngine.emit(.error("need to set SWIFT_EXEC"))
+  exit(1)
+}
+class PrebuitGenDelegate: JobExecutionDelegate {
+  func jobStarted(job: Job, arguments: [String], pid: Int) {
+    diagnosticsEngine.emit(.remark("\(job.moduleName) started"))
+  }
+
+  func jobFinished(job: Job, result: ProcessResult, pid: Int) {
+    diagnosticsEngine.emit(.remark("\(job.moduleName) finished"))
+  }
+
+  func jobSkipped(job: Job) {}
+}
+do {
+  let processSet = ProcessSet()
+  let collector = try SDKPrebuiltModuleInputsCollector(VirtualPath(path: sdkPath).absolutePath!)
+  let inputMap = try collector.collectSwiftInterfaceMap()
+  let allModules = inputMap.keys
+  try withTemporaryFile(suffix: ".swift") {
+    let tempPath = $0.path
+    try localFileSystem.writeFileContents(tempPath, body: {
+      for module in allModules {
+        if module == "MediaPlayer" || module == "PhotosUI" {
+          continue
+        }
+        $0 <<< "import " <<< module <<< "\n"
+      }
+    })
+    print(try localFileSystem.readFileContents(tempPath))
+    let executor = try SwiftDriverExecutor(diagnosticsEngine: diagnosticsEngine,
+                                           processSet: processSet,
+                                           fileSystem: localFileSystem,
+                                           env: ProcessEnv.vars)
+    var driver = try Driver(args: ["swiftc", tempPath.description, "-sdk", sdkPath],
+                            diagnosticsEngine: diagnosticsEngine,
+                            executor: executor,
+                            compilerExecutableDir: VirtualPath.init(path: swiftcPath).parentDirectory.absolutePath!)
+    let jobs = try driver.generatePrebuitModuleGenerationJobs(inputMap,
+      VirtualPath(path: "/tmp/t").absolutePath!)
+    print(jobs.count)
+    try executor.execute(workload: .all(jobs), delegate: PrebuitGenDelegate(), numParallelJobs: 10)
+  }
+} catch {
+  print("error: \(error)")
+  exit(1)
+}

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/A.swiftmodule/arm64e-apple-macos.swiftinterface

@@ -0,0 +1,4 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name A
+import Swift
+public func FuncA() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/A.swiftmodule/x86_64-apple-macos.swiftinterface

@@ -0,0 +1,4 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name A
+import Swift
+public func FuncA() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/E.swiftmodule/arm64e-apple-macos.swiftinterface

@@ -0,0 +1,4 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name E
+import Swift
+public func FuncE() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/E.swiftmodule/x86_64-apple-macos.swiftinterface

@@ -0,0 +1,4 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name E
+import Swift
+public func FuncE() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/F.swiftmodule/arm64e-apple-macos.swiftinterface

@@ -0,0 +1,5 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name F
+import Swift
+import A
+public func FuncF() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/F.swiftmodule/x86_64-apple-macos.swiftinterface

@@ -0,0 +1,5 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name F
+import Swift
+import A
+public func FuncF() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/G.swiftmodule/arm64e-apple-macos.swiftinterface

@@ -0,0 +1,5 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name G
+import Swift
+import E
+public func FuncG() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/G.swiftmodule/x86_64-apple-macos.swiftinterface

@@ -0,0 +1,5 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name G
+import Swift
+import E
+public func FuncG() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/H.swiftmodule/arm64e-apple-macos.swiftinterface

@@ -0,0 +1,8 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name H
+import Swift
+import A
+import E
+import F
+import G
+public func FuncH() { }

TestInputs/ExplicitModuleBuilds/Swift/mock-sdk/usr/lib/swift/H.swiftmodule/x86_64-apple-macos.swiftinterface

@@ -0,0 +1,8 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name H
+import Swift
+import A
+import E
+import F
+import G
+public func FuncH() { }