From 4a46d57f36d7acb25e123e37b5a2bb28dbdb9462 Mon Sep 17 00:00:00 2001
From: Bartosz Polaczyk <polac24@gmail.com>
Date: Sun, 28 May 2023 18:40:51 -0700
Subject: [PATCH] Add -supplementary-output-file-map

---
 .../SwiftFrontend/SwiftFrontendArgInput.swift | 32 ++++++----
 .../Commands/Swiftc/SwiftcContext.swift       |  2 +
 .../Swiftc/SwiftcFilemapInputEditor.swift     | 26 +++++++--
 .../Commands/Swiftc/XCSwiftc.swift            |  4 +-
 .../XCSwiftcFrontendMain.swift                |  6 +-
 .../SwiftcFilemapInputEditorTests.swift       | 58 ++++++++++++++++---
 6 files changed, 101 insertions(+), 27 deletions(-)

diff --git a/Sources/XCRemoteCache/Commands/SwiftFrontend/SwiftFrontendArgInput.swift b/Sources/XCRemoteCache/Commands/SwiftFrontend/SwiftFrontendArgInput.swift
index 26cc2d8c..dda0b23d 100644
--- a/Sources/XCRemoteCache/Commands/SwiftFrontend/SwiftFrontendArgInput.swift
+++ b/Sources/XCRemoteCache/Commands/SwiftFrontend/SwiftFrontendArgInput.swift
@@ -63,6 +63,7 @@ public struct SwiftFrontendArgInput {
     // .swiftsourceinfo and .swiftdoc will be placed next to the .swiftmodule
     let sourceInfoPath: String?
     let docPath: String?
+    let supplementaryOutputFileMap: String?
 
     /// Manual initializer implementation required to be public
     public init(
@@ -77,7 +78,8 @@ public struct SwiftFrontendArgInput {
         dependenciesPaths: [String],
         diagnosticsPaths: [String],
         sourceInfoPath: String?,
-        docPath: String?
+        docPath: String?,
+        supplementaryOutputFileMap: String?
     ) {
         self.compile = compile
         self.emitModule = emitModule
@@ -91,6 +93,7 @@ public struct SwiftFrontendArgInput {
         self.diagnosticsPaths = diagnosticsPaths
         self.sourceInfoPath = sourceInfoPath
         self.docPath = docPath
+        self.supplementaryOutputFileMap = supplementaryOutputFileMap
     }
 
     // swiftlint:disable:next cyclomatic_complexity function_body_length
@@ -139,23 +142,28 @@ public struct SwiftFrontendArgInput {
                 emitModule: nil
             )
 
-            let compilationFileMap = (0..<primaryInputsCount).reduce([String: SwiftFileCompilationInfo]()) { prev, i in
-                var new = prev
-                new[primaryInputPaths[i]] = SwiftFileCompilationInfo(
-                    file: primaryInputFilesURLs[i],
-                    dependencies: dependenciesPaths.get(i).map(URL.init(fileURLWithPath:)),
-                    object: outputPaths.get(i).map(URL.init(fileURLWithPath:)),
-                    // for now - swift-dependencies are not requested in the driver compilation mode
-                    swiftDependencies: nil
-                )
-                return new
+            let compilationFilesOutputs: SwiftcContext.CompilationFilesOutputs
+            if let compimentaryFileMa = supplementaryOutputFileMap {
+                compilationFilesOutputs = .supplementaryFileMap(compimentaryFileMa)
+            } else {
+                compilationFilesOutputs = .map((0..<primaryInputsCount).reduce([String: SwiftFileCompilationInfo]()) { prev, i in
+                    var new = prev
+                    new[primaryInputPaths[i]] = SwiftFileCompilationInfo(
+                        file: primaryInputFilesURLs[i],
+                        dependencies: dependenciesPaths.get(i).map(URL.init(fileURLWithPath:)),
+                        object: outputPaths.get(i).map(URL.init(fileURLWithPath:)),
+                        // for now - swift-dependencies are not requested in the driver compilation mode
+                        swiftDependencies: nil
+                    )
+                    return new
+                })
             }
 
             return try .init(
                 config: config,
                 moduleName: moduleName,
                 steps: steps,
-                outputs: .map(compilationFileMap),
+                outputs: compilationFilesOutputs,
                 target: target,
                 inputs: .list(inputPaths),
                 exampleWorkspaceFilePath: outputPaths[0]
diff --git a/Sources/XCRemoteCache/Commands/Swiftc/SwiftcContext.swift b/Sources/XCRemoteCache/Commands/Swiftc/SwiftcContext.swift
index 03b9f20e..47c2a964 100644
--- a/Sources/XCRemoteCache/Commands/Swiftc/SwiftcContext.swift
+++ b/Sources/XCRemoteCache/Commands/Swiftc/SwiftcContext.swift
@@ -48,6 +48,8 @@ public struct SwiftcContext {
     public enum CompilationFilesOutputs {
         /// defined in a separate file (via -output-file-map)
         case fileMap(String)
+        /// defined in a separate file (via -supplementary-output-file-map)
+        case supplementaryFileMap(String)
         /// explicitly passed in the invocation
         case map([String: SwiftFileCompilationInfo])
     }
diff --git a/Sources/XCRemoteCache/Commands/Swiftc/SwiftcFilemapInputEditor.swift b/Sources/XCRemoteCache/Commands/Swiftc/SwiftcFilemapInputEditor.swift
index 0d8776ed..d4c23cd3 100644
--- a/Sources/XCRemoteCache/Commands/Swiftc/SwiftcFilemapInputEditor.swift
+++ b/Sources/XCRemoteCache/Commands/Swiftc/SwiftcFilemapInputEditor.swift
@@ -18,11 +18,13 @@
 // under the License.
 
 import Foundation
+import Yams
 
 /// Errors with reading swiftc inputs
 enum SwiftcInputReaderError: Error {
     case readingFailed
     case invalidFormat
+    case invalidYamlFormat
     case missingField(String)
 }
 
@@ -49,7 +51,7 @@ struct SwiftModuleCompilationInfo: Encodable, Equatable {
     let swiftDependencies: URL?
 }
 
-public struct SwiftFileCompilationInfo: Encodable, Equatable {
+public struct SwiftFileCompilationInfo: Encodable, Hashable {
     let file: URL
     // not present for WMO builds
     let dependencies: URL?
@@ -61,11 +63,18 @@ public struct SwiftFileCompilationInfo: Encodable, Equatable {
 
 class SwiftcFilemapInputEditor: SwiftcInputReader, SwiftcInputWriter {
 
+    enum Format {
+        case json
+        case yaml
+    }
+
     private let file: URL
+    private let fileFormat: Format
     private let fileManager: FileManager
 
-    init(_ file: URL, fileManager: FileManager) {
+    init(_ file: URL, fileFormat: Format, fileManager: FileManager) {
         self.file = file
+        self.fileFormat = fileFormat
         self.fileManager = fileManager
     }
 
@@ -73,7 +82,7 @@ class SwiftcFilemapInputEditor: SwiftcInputReader, SwiftcInputWriter {
         guard let content = fileManager.contents(atPath: file.path) else {
             throw SwiftcInputReaderError.readingFailed
         }
-        guard let representation = try JSONSerialization.jsonObject(with: content, options: []) as? [String: Any] else {
+        guard let representation = try decodeFile(content: content) else {
             throw SwiftcInputReaderError.invalidFormat
         }
         return try SwiftCompilationInfo(from: representation)
@@ -83,11 +92,20 @@ class SwiftcFilemapInputEditor: SwiftcInputReader, SwiftcInputWriter {
         let data = try JSONSerialization.data(withJSONObject: info.dump(), options: [.prettyPrinted])
         fileManager.createFile(atPath: file.path, contents: data, attributes: nil)
     }
+
+    private func decodeFile(content: Data) throws -> [String: Any]? {
+        switch fileFormat {
+        case .json:
+            return try JSONSerialization.jsonObject(with: content, options: []) as? [String: Any]
+        case .yaml:
+            return try Yams.load(yaml: String(data: content, encoding: .utf8)!) as? [String: Any]
+        }
+    }
 }
 
 extension SwiftCompilationInfo {
     init(from object: [String: Any]) throws {
-        info = try SwiftModuleCompilationInfo(from: object[""])
+        info = try SwiftModuleCompilationInfo(from: object["", default: [:]])
         files = try object.reduce([]) { prev, new in
             let (key, value) = new
             if key.isEmpty {
diff --git a/Sources/XCRemoteCache/Commands/Swiftc/XCSwiftc.swift b/Sources/XCRemoteCache/Commands/Swiftc/XCSwiftc.swift
index ec2a54b0..a071f93c 100644
--- a/Sources/XCRemoteCache/Commands/Swiftc/XCSwiftc.swift
+++ b/Sources/XCRemoteCache/Commands/Swiftc/XCSwiftc.swift
@@ -83,7 +83,9 @@ public class XCSwiftAbstract<InputArgs> {
         // and the file/list is named "output"
         switch context.outputs {
         case .fileMap(let path):
-            inputReader = SwiftcFilemapInputEditor(URL(fileURLWithPath: path), fileManager: fileManager)
+            inputReader = SwiftcFilemapInputEditor(URL(fileURLWithPath: path), fileFormat: .json, fileManager: fileManager)
+        case .supplementaryFileMap(let path):
+            inputReader = SwiftcFilemapInputEditor(URL(fileURLWithPath: path), fileFormat: .yaml, fileManager: fileManager)
         case .map(let map):
             // static - passed via the arguments list
             // TODO: check if first 2 ars can always be `nil`
diff --git a/Sources/xcswift-frontend/XCSwiftcFrontendMain.swift b/Sources/xcswift-frontend/XCSwiftcFrontendMain.swift
index 2e730fdb..05141f49 100644
--- a/Sources/xcswift-frontend/XCSwiftcFrontendMain.swift
+++ b/Sources/xcswift-frontend/XCSwiftcFrontendMain.swift
@@ -42,6 +42,7 @@ public class XCSwiftcFrontendMain {
         var diagnosticsPaths: [String] = []
         var sourceInfoPath: String?
         var docPath: String?
+        var supplementaryOutputFileMap: String?
 
         for i in 0..<args.count {
             let arg = args[i]
@@ -73,6 +74,8 @@ public class XCSwiftcFrontendMain {
             case "-primary-file":
                 // .swift
                 primaryInputPaths.append(args[i + 1])
+            case "-supplementary-output-file-map":
+                supplementaryOutputFileMap = args[i + 1]
             default:
                 if arg.hasSuffix(".swift") {
                     inputPaths.append(arg)
@@ -94,7 +97,8 @@ public class XCSwiftcFrontendMain {
             dependenciesPaths: dependenciesPaths,
             diagnosticsPaths: diagnosticsPaths,
             sourceInfoPath: sourceInfoPath,
-            docPath: docPath
+            docPath: docPath,
+            supplementaryOutputFileMap: supplementaryOutputFileMap
         )
         // swift-frontened is first invoked with some "probing" args like
         // -print-target-info
diff --git a/Tests/XCRemoteCacheTests/Commands/SwiftcFilemapInputEditorTests.swift b/Tests/XCRemoteCacheTests/Commands/SwiftcFilemapInputEditorTests.swift
index 80b2eb15..8bd688ba 100644
--- a/Tests/XCRemoteCacheTests/Commands/SwiftcFilemapInputEditorTests.swift
+++ b/Tests/XCRemoteCacheTests/Commands/SwiftcFilemapInputEditorTests.swift
@@ -34,19 +34,21 @@ class SwiftcFilemapInputEditorTests: FileXCTestCase {
     )
     private let sampleInfoContentData = #"{"":{"swift-dependencies":"/"}}"#.data(using: .utf8)!
     private var inputFile: URL!
-    private var editor: SwiftcFilemapInputEditor!
+    private var editorJson: SwiftcFilemapInputEditor!
+    private var editorYaml: SwiftcFilemapInputEditor!
 
     override func setUpWithError() throws {
         try super.setUpWithError()
         try prepareTempDir()
         inputFile = workingDirectory!.appendingPathComponent("swift.json")
-        editor = SwiftcFilemapInputEditor(inputFile, fileManager: fileManager)
+        editorJson = SwiftcFilemapInputEditor(inputFile, fileFormat: .json, fileManager: fileManager)
+        editorYaml = SwiftcFilemapInputEditor(inputFile, fileFormat: .yaml, fileManager: fileManager)
     }
 
     func testReading() throws {
         try fileManager.spt_writeToFile(atPath: inputFile.path, contents: sampleInfoContentData)
 
-        let readInfo = try editor.read()
+        let readInfo = try editorJson.read()
 
         XCTAssertEqual(readInfo, sampleInfo)
     }
@@ -80,13 +82,13 @@ class SwiftcFilemapInputEditorTests: FileXCTestCase {
             ])
         try fileManager.spt_writeToFile(atPath: inputFile.path, contents: infoContentData)
 
-        let readInfo = try editor.read()
+        let readInfo = try editorJson.read()
 
         XCTAssertEqual(readInfo, expectedInfo)
     }
 
     func testWritingSavesContent() throws {
-        try editor.write(sampleInfo)
+        try editorJson.write(sampleInfo)
 
         let savedContent = try Data(contentsOf: inputFile)
         let content = try JSONSerialization.jsonObject(with: savedContent, options: []) as? [String: Any]
@@ -108,7 +110,7 @@ class SwiftcFilemapInputEditorTests: FileXCTestCase {
                 ),
             ])
 
-        try editor.write(extendedInfo)
+        try editorJson.write(extendedInfo)
 
         let savedContent = try Data(contentsOf: inputFile)
         let content = try JSONSerialization.jsonObject(with: savedContent, options: []) as? [String: Any]
@@ -119,12 +121,50 @@ class SwiftcFilemapInputEditorTests: FileXCTestCase {
     func testModifyingFileCompilationInfo() throws {
         try fileManager.spt_writeToFile(atPath: inputFile.path, contents: sampleInfoContentData)
 
-        let originalInfo = try editor.read()
+        let originalInfo = try editorJson.read()
         var modifiedInfo = originalInfo
         modifiedInfo.files = [file]
-        try editor.write(modifiedInfo)
-        let finalInfo = try editor.read()
+        try editorJson.write(modifiedInfo)
+        let finalInfo = try editorJson.read()
 
         XCTAssertEqual(finalInfo, modifiedInfo)
     }
+
+    func testReadingSupplementaryInfoWithOptionalProperties() throws {
+        let infoContentData = #"""
+        "/file1.swift":
+          swift-dependencies: "/file1.swiftdeps"
+          dependencies: "/file1.d"
+        "/file2.swift":
+          dependencies: "/file2.d"
+          object: "/file2.o"
+          swift-dependencies: "/file2.swiftdeps"
+        """#.data(using: .utf8)!
+        let expectedInfo = SwiftCompilationInfo(
+            info: SwiftModuleCompilationInfo(
+                dependencies: nil,
+                swiftDependencies: nil
+            ),
+            files: [
+                SwiftFileCompilationInfo(
+                    file: "/file1.swift",
+                    dependencies: "/file1.d",
+                    object: nil,
+                    swiftDependencies: "/file1.swiftdeps"
+                ),
+                SwiftFileCompilationInfo(
+                    file: "/file2.swift",
+                    dependencies: "/file2.d",
+                    object: "/file2.o",
+                    swiftDependencies: "/file2.swiftdeps"
+                ),
+            ])
+        try fileManager.spt_writeToFile(atPath: inputFile.path, contents: infoContentData)
+
+        let readInfo = try editorYaml.read()
+
+        // `files` order doesn't match
+        XCTAssertEqual(readInfo.info, expectedInfo.info)
+        XCTAssertEqual(Set(readInfo.files), Set(expectedInfo.files))
+    }
 }