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/FrontendJobHelpers.swift

@@ -79,12 +79,6 @@ extension Driver {
       commandLine.appendFlag(Triple(variant, normalizing: true).triple)
     }
 
-    // Enable address top-byte ignored in the ARM64 backend.
-    if targetTriple.arch == .aarch64 {
-      commandLine.appendFlag(.Xllvm)
-      commandLine.appendFlag("-aarch64-use-tbi")
-    }
-
     // Enable or disable ObjC interop appropriately for the platform
     if targetTriple.isDarwin {
       commandLine.appendFlag(.enableObjcInterop)

Sources/SwiftDriver/Jobs/Planning.swift

@@ -496,6 +496,110 @@ extension Driver {
     return try explicitDependencyBuildPlanner!.generateExplicitModuleDependenciesBuildJobs()
   }
 
+
+  public mutating func generatePrebuitModuleGenerationJobs(_ inputMap: [String: [TypedVirtualPath]],
+                                                           _ prebuiltModuleDir: AbsolutePath) throws -> [Job] {
+    guard let rawSdkPath = frontendTargetInfo.sdkPath?.path else {
+      // cannot proceed without SDK path
+      return []
+    }
+    let sdkPath = VirtualPath.lookup(rawSdkPath)
+    // 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 {
+      assert(!$0.value.isEmpty)
+      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 isIosMac(_ path: TypedVirtualPath) -> Bool {
+      return path.file.basenameWithoutExt.contains("macabi")
+    }
+
+    func getDependenciesPaths(_ module: String, _ iosMac: Bool) 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.filter { isIosMac($0) == iosMac })
+          }
+        }
+      }
+      return results
+    }
+
+    let iosMacFrameworks = sdkPath
+      .appending(component: "System")
+      .appending(component: "iOSSupport")
+      .appending(component: "System")
+      .appending(component: "Library")
+      .appending(component: "Frameworks")
+    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)
+          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(sdkPath))
+            commandLine.appendFlag(.prebuiltModuleCachePath)
+            commandLine.appendPath(prebuiltModuleDir)
+            commandLine.appendFlag(.moduleName)
+            commandLine.appendFlag(moduleName)
+            commandLine.appendFlag(.o)
+            commandLine.appendPath(outputPath.file)
+            commandLine.appendPath(inputPath.file)
+            if moduleName == "Swift" {
+              commandLine.appendFlag(.parseStdlib)
+            }
+            if isIosMac(inputPath) {
+              commandLine.appendFlag(.Fsystem)
+              commandLine.append(.path(iosMacFrameworks))
+            }
+            commandLine.appendFlag(.serializeParseableModuleInterfaceDependencyHashes)
+            jobs.append(Job(
+              moduleName: moduleName,
+              kind: .compile,
+              tool: .absolute(try toolchain.getToolPath(.swiftCompiler)),
+              commandLine: commandLine,
+              inputs: try getDependenciesPaths(moduleName, isIosMac(inputPath)),
+              primaryInputs: [],
+              outputs: [outputPath]
+            ))
+          }
+        }
+      }
+    }
+
+    return jobs
+  }
+
   private mutating func gatherModuleDependencies()
   throws -> InterModuleDependencyGraph {
     var dependencyGraph = try performDependencyScan()
@@ -551,6 +655,107 @@ 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")]
+  let sdkInfo: DarwinToolchain.DarwinSDKInfo
+  let diagEngine: DiagnosticsEngine
+  public init(_ sdkPath: AbsolutePath, _ diagEngine: DiagnosticsEngine) {
+    self.sdkPath = sdkPath
+    self.sdkInfo = DarwinToolchain.readSDKInfo(localFileSystem,
+                                               VirtualPath.absolute(sdkPath).intern())!
+    self.diagEngine = diagEngine
+  }
+
+  public var targetTriple: String {
+    let canonicalName = sdkInfo.canonicalName
+    func extractVersion(_ platform: String) -> Substring? {
+      if canonicalName.starts(with: platform) {
+        return canonicalName.suffix(from: canonicalName.index(canonicalName.startIndex,
+                                                              offsetBy: platform.count))
+      }
+      return nil
+    }
+
+    if let version = extractVersion("macosx") {
+      return "x86_64-apple-macosx\(version)"
+    } else if let version = extractVersion("iphoneos") {
+      return "arm64-apple-ios\(version)"
+    } else {
+      diagEngine.emit(error: "unhandled platform name: \(canonicalName)")
+      return ""
+    }
+  }
+
+  private func sanitizeInterfaceMap(_ map: [String: [TypedVirtualPath]]) -> [String: [TypedVirtualPath]] {
+    return map.filter {
+      if $0.value.isEmpty {
+        diagEngine.emit(warning: "\($0.key) has no associated .swiftinterface files")
+        return false
+      }
+      return true
+    }
+  }
+
+  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 = AbsolutePath(dir, try VirtualPath(path: $0).relativePath!)
+        if currentFile.extension == "swiftinterface" {
+          let currentBaseName = currentFile.basenameWithoutExt
+          let interfacePath = TypedVirtualPath(file: VirtualPath.absolute(currentFile).intern(),
+                                               type: .swiftInterface)
+          if !results[moduleName]!.contains(where: { $0.file.basenameWithoutExt == currentBaseName }) {
+            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 {
+        if $0.hasSuffix(".swiftmodule") {
+          try updateResults(AbsolutePath(swiftModuleDir, $0))
+        }
+      }
+    }
+    return sanitizeInterfaceMap(results)
+  }
+}
+
 /// MARK: Planning
 extension Driver {
   /// Create a job if needed for simple requests that can be immediately

Sources/SwiftDriver/Toolchains/DarwinToolchain.swift

@@ -227,15 +227,17 @@ import SwiftOptions
     private enum CodingKeys: String, CodingKey {
       case version = "Version"
       case versionMap = "VersionMap"
+      case canonicalName = "CanonicalName"
     }
 
     struct VersionMap: Decodable {
       private enum CodingKeys: String, CodingKey {
         case macOSToCatalystMapping = "macOS_iOSMac"
       }
 
-      var macOSToCatalystMapping: [Version: Version]
+      var macOSToCatalystMapping: [Version: Version] = [:]
 
+      init() {}
       init(from decoder: Decoder) throws {
         let keyedContainer = try decoder.container(keyedBy: CodingKeys.self)
 
@@ -259,18 +261,23 @@ import SwiftOptions
 
     private var version: Version
     private var versionMap: VersionMap
-
+    let canonicalName: String
     init(from decoder: Decoder) throws {
       let keyedContainer = try decoder.container(keyedBy: CodingKeys.self)
 
       let versionString = try keyedContainer.decode(String.self, forKey: .version)
+      self.canonicalName = try keyedContainer.decode(String.self, forKey: .canonicalName)
       guard let version = Version(potentiallyIncompleteVersionString: versionString) else {
         throw DecodingError.dataCorruptedError(forKey: .version,
                                                in: keyedContainer,
                                                debugDescription: "Malformed version string")
       }
       self.version = version
-      self.versionMap = try keyedContainer.decode(VersionMap.self, forKey: .versionMap)
+      if self.canonicalName.hasPrefix("macosx") {
+        self.versionMap = try keyedContainer.decode(VersionMap.self, forKey: .versionMap)
+      } else {
+        self.versionMap = VersionMap()
+      }
     }
 
     func sdkVersion(for triple: Triple) -> Version {
@@ -287,16 +294,20 @@ import SwiftOptions
   // SDK info is computed lazily. This should not generally be accessed directly.
   private var _sdkInfo: DarwinSDKInfo? = nil
 
+  static func readSDKInfo(_ fileSystem: FileSystem, _ sdkPath: VirtualPath.Handle) -> DarwinSDKInfo? {
+    let sdkSettingsPath = VirtualPath.lookup(sdkPath).appending(component: "SDKSettings.json")
+    guard let contents = try? fileSystem.readFileContents(sdkSettingsPath) else { return nil }
+    guard let sdkInfo = try? JSONDecoder().decode(DarwinSDKInfo.self,
+                                                  from: Data(contents.contents)) else { return nil }
+    return sdkInfo
+  }
+
   func getTargetSDKInfo(sdkPath: VirtualPath.Handle) -> DarwinSDKInfo? {
     if let info = _sdkInfo {
       return info
     } else {
-      let sdkSettingsPath = VirtualPath.lookup(sdkPath).appending(component: "SDKSettings.json")
-      guard let contents = try? fileSystem.readFileContents(sdkSettingsPath) else { return nil }
-      guard let sdkInfo = try? JSONDecoder().decode(DarwinSDKInfo.self,
-                                                    from: Data(contents.contents)) else { return nil }
-      self._sdkInfo = sdkInfo
-      return sdkInfo
+      self._sdkInfo = DarwinToolchain.readSDKInfo(fileSystem, sdkPath)
+      return self._sdkInfo
     }
   }
 

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

@@ -0,0 +1,72 @@
+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 {
+  var commandMap: [Int: String] = [:]
+  func jobStarted(job: Job, arguments: [String], pid: Int) {
+    commandMap[pid] = arguments.reduce("") { return $0 + " " + $1 }
+  }
+
+  func jobFinished(job: Job, result: ProcessResult, pid: Int) {
+    switch result.exitStatus {
+    case .terminated(code: let code):
+      if code != 0 {
+        Driver.stdErrQueue.sync {
+          stderrStream <<< "failed: " <<< commandMap[pid]! <<< "\n"
+          stderrStream.flush()
+        }
+      }
+    case .signalled:
+      diagnosticsEngine.emit(.remark("\(job.moduleName) interrupted"))
+    }
+  }
+
+  func jobSkipped(job: Job) {
+    diagnosticsEngine.emit(.error("\(job.moduleName) skipped"))
+  }
+}
+do {
+  let processSet = ProcessSet()
+  let collector = try SDKPrebuiltModuleInputsCollector(VirtualPath(path: sdkPath).absolutePath!, diagnosticsEngine)
+  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 {
+        $0 <<< "import " <<< module <<< "\n"
+      }
+    })
+    let executor = try SwiftDriverExecutor(diagnosticsEngine: diagnosticsEngine,
+                                           processSet: processSet,
+                                           fileSystem: localFileSystem,
+                                           env: ProcessEnv.vars)
+    var driver = try Driver(args: ["swiftc",
+                                   "-target", collector.targetTriple,
+                                   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!)
+
+    try executor.execute(workload: DriverExecutorWorkload.init(jobs, nil, continueBuildingAfterErrors: true), delegate: PrebuitGenDelegate(), numParallelJobs: 1024)
+  }
+} catch {
+  print("error: \(error)")
+  exit(1)
+}