From 7da6a5e5a54d9c0615d1134b0448c1c153f1a372 Mon Sep 17 00:00:00 2001
From: Xi Ge <xi_ge@apple.com>
Date: Thu, 10 Jun 2021 16:41:17 -0700
Subject: [PATCH] PrebuiltModuleGen: teach the tool to dump a .dot file for
 module dependency visualization

---
 .../SwiftDriver/Jobs/PrebuiltModulesJob.swift | 57 ++++++++++++++++++-
 Sources/swift-build-sdk-interfaces/main.swift | 10 +++-
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftDriver/Jobs/PrebuiltModulesJob.swift b/Sources/SwiftDriver/Jobs/PrebuiltModulesJob.swift
index a57dad636..7e33921d1 100644
--- a/Sources/SwiftDriver/Jobs/PrebuiltModulesJob.swift
+++ b/Sources/SwiftDriver/Jobs/PrebuiltModulesJob.swift
@@ -239,6 +239,57 @@ public struct SDKPrebuiltModuleInputsCollector {
   }
 }
 
+extension InterModuleDependencyGraph {
+  func dumpDotGraph(_ path: AbsolutePath, _ includingPCM: Bool) throws {
+    func isPCM(_ dep: ModuleDependencyId) -> Bool {
+      switch dep {
+      case .clang:
+        return true
+      default:
+        return false
+      }
+    }
+    func dumpModuleName(_ stream: WritableByteStream, _ dep: ModuleDependencyId) {
+      switch dep {
+      case .swift(let name):
+        stream <<< "\"\(name).swiftmodule\""
+      case .clang(let name):
+        stream <<< "\"\(name).pcm\""
+      default:
+        break
+      }
+    }
+    try localFileSystem.writeFileContents(path) {Stream in
+      Stream <<< "digraph {\n"
+      for key in modules.keys {
+        switch key {
+        case .swift(let name):
+          if name == mainModuleName {
+              break
+          }
+          fallthrough
+        case .clang:
+          if !includingPCM && isPCM(key) {
+            break
+          }
+          modules[key]!.directDependencies?.forEach { dep in
+            if !includingPCM && isPCM(dep) {
+              return
+            }
+            dumpModuleName(Stream, key)
+            Stream <<< " -> "
+            dumpModuleName(Stream, dep)
+            Stream <<< ";\n"
+          }
+        default:
+          break
+        }
+      }
+      Stream <<< "}\n"
+    }
+  }
+}
+
 extension Driver {
 
   private mutating func generateSingleModuleBuildingJob(_ moduleName: String,  _ prebuiltModuleDir: AbsolutePath,
@@ -285,12 +336,16 @@ extension Driver {
 
   public mutating func generatePrebuitModuleGenerationJobs(with inputMap: [String: [PrebuiltModuleInput]],
                                                            into prebuiltModuleDir: AbsolutePath,
-                                                           exhaustive: Bool) throws -> ([Job], [Job]) {
+                                                           exhaustive: Bool,
+                                                           dotGraphPath: AbsolutePath? = nil) throws -> ([Job], [Job]) {
     assert(sdkPath != nil)
     // Run the dependency scanner and update the dependency oracle with the results
     // We only need Swift dependencies here, so we don't need to invoke gatherModuleDependencies,
     // which also resolves versioned clang modules.
     let dependencyGraph = try performDependencyScan()
+    if let dotGraphPath = dotGraphPath {
+      try dependencyGraph.dumpDotGraph(dotGraphPath, false)
+    }
     var jobs: [Job] = []
     var danglingJobs: [Job] = []
     var inputCount = 0
diff --git a/Sources/swift-build-sdk-interfaces/main.swift b/Sources/swift-build-sdk-interfaces/main.swift
index 3b72f0b8a..6cb8418df 100644
--- a/Sources/swift-build-sdk-interfaces/main.swift
+++ b/Sources/swift-build-sdk-interfaces/main.swift
@@ -32,6 +32,13 @@ func getArgument(_ flag: String) -> String? {
   return nil
 }
 
+func getArgumentAsPath(_ flag: String) throws -> AbsolutePath? {
+  if let raw = getArgument(flag) {
+    return try VirtualPath(path: raw).absolutePath
+  }
+  return nil
+}
+
 guard let rawOutputDir = getArgument("-o") else {
   diagnosticsEngine.emit(.error("need to specify -o"))
   exit(1)
@@ -119,7 +126,8 @@ do {
                             diagnosticsEngine: diagnosticsEngine,
                             executor: executor,
                             compilerExecutableDir: swiftcPath.parentDirectory)
-    let (jobs, danglingJobs) = try driver.generatePrebuitModuleGenerationJobs(with: inputMap, into: outputDir, exhaustive: !coreMode)
+    let (jobs, danglingJobs) = try driver.generatePrebuitModuleGenerationJobs(with: inputMap,
+      into: outputDir, exhaustive: !coreMode, dotGraphPath: getArgumentAsPath("-dot-graph-path"))
     if verbose {
       Driver.stdErrQueue.sync {
         stderrStream <<< "job count: \(jobs.count + danglingJobs.count)\n"