[Macros] Add support for wasm plugins

Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift

@@ -318,7 +318,10 @@ extension Driver {
     }
 
     // Emit user-provided plugin paths, in order.
-    if isFrontendArgSupported(.externalPluginPath) {
+    if isFrontendArgSupported(.loadPlugin) {
+      let options = parsedOptions.arguments(for: .pluginPath, .externalPluginPath, .loadPluginLibrary, .loadPluginExecutable, .loadPlugin)
+      try commandLine.append(contentsOf: handleWasmPlugins(options))
+    } else if isFrontendArgSupported(.externalPluginPath) {
       try commandLine.appendAll(.pluginPath, .externalPluginPath, .loadPluginLibrary, .loadPluginExecutable, from: &parsedOptions)
     } else if isFrontendArgSupported(.pluginPath) {
       try commandLine.appendAll(.pluginPath, .loadPluginLibrary, from: &parsedOptions)
@@ -818,6 +821,21 @@ extension Driver {
     commandLine.appendPath(pluginPathRoot.localPluginPath)
   }
 
+  /// Forward all plugin-related arguments, with special handling for Wasm plugins requested via `-load-plugin-executable`.
+  ///
+  /// Wasm plugins are updated to use `-load-plugin` with the plugin server.
+  private func handleWasmPlugins(_ options: [ParsedOption]) throws -> [ParsedOption] {
+    lazy var pluginServerPathResult = Result {
+      VirtualPath.absolute(try toolchain.executableDir.parentDirectory).pluginServerPath
+    }
+    return try options.map { option in
+      try option.handlingWasmPlugins(
+        // Not the same as `pluginServerPath: pluginServerPathResult.get`.
+        // We want to ensure that `pluginServerPathResult` is lazily initialized.
+        pluginServerPath: { try pluginServerPathResult.get() }
+      )
+    }
+  }
 
   /// If explicit dependency planner supports creating bridging header pch command.
   public func supportsBridgingHeaderPCHCommand() throws -> Bool {
@@ -951,3 +969,36 @@ extension ParsedOptions {
     }
   }
 }
+
+extension ParsedOption {
+  /// Transforms the receiver to handle Wasm plugins.
+  ///
+  /// If the option is of the form `-load-plugin-executable foo.wasm`, it's updated to use
+  /// `-load-plugin` with the plugin server. Otherwise, it's returned unmodified.
+  fileprivate func handlingWasmPlugins(
+    pluginServerPath: () throws -> VirtualPath
+  ) throws -> ParsedOption {
+    guard option == .loadPluginExecutable else {
+      return self
+    }
+
+    let argument = argument.asSingle
+    guard let separator = argument.lastIndex(of: "#") else {
+      return self
+    }
+
+    let path = argument[..<separator]
+    guard path.hasSuffix(".wasm") else {
+      return self
+    }
+    let afterPath = argument[separator...]
+
+    let pluginServerPath = try pluginServerPath()
+
+    return ParsedOption(
+      option: .loadPlugin,
+      argument: .single("\(path):\(pluginServerPath)\(afterPath)"),
+      index: index
+    )
+  }
+}

Sources/SwiftDriver/Toolchains/DarwinToolchain.swift

@@ -501,16 +501,6 @@ private extension Version {
 }
 
 extension VirtualPath {
-  // Given a virtual path pointing into a toolchain/SDK/platform, produce the
-  // path to `swift-plugin-server`.
-  fileprivate var pluginServerPath: VirtualPath {
-#if os(Windows)
-    self.appending(components: "bin", "swift-plugin-server.exe")
-#else
-    self.appending(components: "bin", "swift-plugin-server")
-#endif
-  }
-
   // Given a virtual path pointing into a toolchain/SDK/platform, produce the
   // path to the plugins.
   var pluginPath: VirtualPath {

Sources/SwiftDriver/Toolchains/Toolchain.swift

@@ -371,3 +371,15 @@ extension Toolchain {
 @_spi(Testing) public enum ToolchainError: Swift.Error {
   case unableToFind(tool: String)
 }
+
+extension VirtualPath {
+  // Given a virtual path pointing into a toolchain/SDK/platform, produce the
+  // path to `swift-plugin-server`.
+  var pluginServerPath: VirtualPath {
+#if os(Windows)
+    self.appending(components: "bin", "swift-plugin-server.exe")
+#else
+    self.appending(components: "bin", "swift-plugin-server")
+#endif
+  }
+}

Sources/SwiftOptions/Options.swift

@@ -594,6 +594,7 @@ extension Option {
   public static let reuseDependencyScanCache: Option = Option("-load-dependency-scan-cache", .flag, attributes: [.frontend, .noDriver], helpText: "After performing a dependency scan, serialize the scanner's internal state.")
   public static let loadPluginExecutable: Option = Option("-load-plugin-executable", .separate, attributes: [.frontend, .doesNotAffectIncrementalBuild, .argumentIsPath], metaVar: "<path>#<module-names>", helpText: "Path to a compiler plugin executable and a comma-separated list of module names where the macro types are declared", group: .pluginSearch)
   public static let loadPluginLibrary: Option = Option("-load-plugin-library", .separate, attributes: [.frontend, .doesNotAffectIncrementalBuild, .argumentIsPath], metaVar: "<path>", helpText: "Path to a dynamic library containing compiler plugins such as macros", group: .pluginSearch)
+  public static let loadPlugin: Option = Option("-load-plugin", .separate, attributes: [.frontend, .doesNotAffectIncrementalBuild], metaVar: "<library-path>:<server-path>#<module-names>", helpText: "Path to a plugin library, a server to load it in, and a comma-separated list of module names where the macro types are declared", group: .pluginSearch)
   public static let locale: Option = Option("-locale", .separate, attributes: [.frontend, .doesNotAffectIncrementalBuild], metaVar: "<locale-code>", helpText: "Choose a language for diagnostic messages")
   public static let localizationPath: Option = Option("-localization-path", .separate, attributes: [.frontend, .doesNotAffectIncrementalBuild, .argumentIsPath], metaVar: "<path>", helpText: "Path to localized diagnostic messages directory")
   public static let location: Option = Option("-location", .separate, attributes: [.noDriver], metaVar: "<location>", helpText: "Filter nodes with the given location.")
@@ -1470,6 +1471,7 @@ extension Option {
       Option.reuseDependencyScanCache,
       Option.loadPluginExecutable,
       Option.loadPluginLibrary,
+      Option.loadPlugin,
       Option.locale,
       Option.localizationPath,
       Option.location,

Tests/SwiftDriverTests/SwiftDriverTests.swift

@@ -7642,7 +7642,7 @@ final class SwiftDriverTests: XCTestCase {
   func pluginPathTest(platform: String, sdk: String, searchPlatform: String) throws {
     let sdkRoot = try testInputsPath.appending(
       components: ["Platform Checks", "\(platform).platform", "Developer", "SDKs", "\(sdk).sdk"])
-    var driver = try Driver(args: ["swiftc", "-typecheck", "foo.swift", "-sdk", VirtualPath.absolute(sdkRoot).name, "-plugin-path", "PluginA", "-external-plugin-path", "Plugin~B#Bexe", "-load-plugin-library", "PluginB2", "-plugin-path", "PluginC", "-working-directory", "/tmp"])
+    var driver = try Driver(args: ["swiftc", "-typecheck", "foo.swift", "-sdk", VirtualPath.absolute(sdkRoot).name, "-plugin-path", "PluginA", "-external-plugin-path", "Plugin~B#Bexe", "-load-plugin-library", "PluginB2", "-plugin-path", "PluginC", "-load-plugin-executable", "PluginD#PluginD", "-working-directory", "/tmp"])
     guard driver.isFrontendArgSupported(.pluginPath) && driver.isFrontendArgSupported(.externalPluginPath) else {
       return
     }
@@ -7698,6 +7698,7 @@ final class SwiftDriverTests: XCTestCase {
     #endif
 
     XCTAssertTrue(job.commandLine.contains(.flag("-plugin-path")))
+    XCTAssertFalse(job.commandLine.contains(.flag("-load-plugin")))
 #if os(Windows)
     XCTAssertTrue(job.commandLine.contains(.path(.absolute(try driver.toolchain.executableDir.parentDirectory.appending(components: "bin")))))
 #else
@@ -7706,6 +7707,42 @@ final class SwiftDriverTests: XCTestCase {
 #endif
   }
 
+  func testWasmPlugins() throws {
+    let sdkRoot = try testInputsPath.appending(
+      components: ["Platform Checks", "iPhoneOS.platform", "Developer", "SDKs", "iPhoneOS13.0.sdk"])
+    let executor = try SwiftDriverExecutor(
+      diagnosticsEngine: DiagnosticsEngine(handlers: [Driver.stderrDiagnosticsHandler]),
+      processSet: ProcessSet(),
+      fileSystem: localFileSystem,
+      env: [:]
+    )
+    let executableDir = try AbsolutePath(validating: "/tmp/swift/bin")
+    var driver = try Driver(
+      args: ["swiftc", "-typecheck", "foo.swift", "-sdk", VirtualPath.absolute(sdkRoot).name, "-load-plugin-executable", "PluginA#ModuleA", "-load-plugin-executable", "PluginB.wasm#ModuleB", "-working-directory", "/tmp", "-load-plugin-library", "PluginC.dylib"],
+      executor: executor,
+      compilerExecutableDir: executableDir
+    )
+    try XCTSkipUnless(driver.isFrontendArgSupported(.loadPlugin))
+
+    let jobs = try driver.planBuild().removingAutolinkExtractJobs()
+    XCTAssertEqual(jobs.count, 1)
+    let job = jobs.first!
+
+    #if os(Windows)
+    let expectedWasmServerPath = executableDir.appending(component: "swift-plugin-server.exe")
+    #else
+    let expectedWasmServerPath = executableDir.appending(component: "swift-plugin-server")
+    #endif
+
+    // Check that the we have the plugin path
+    let pluginAIndex = try XCTUnwrap(job.commandLine.firstIndex(of: .path(.absolute(AbsolutePath(validating: "/tmp/PluginA#ModuleA")))))
+    let pluginBIndex = try XCTUnwrap(job.commandLine.firstIndex(of: .flag("/tmp/PluginB.wasm:\(expectedWasmServerPath)#ModuleB")))
+    let pluginCIndex = try XCTUnwrap(job.commandLine.firstIndex(of: .path(.absolute(AbsolutePath(validating: "/tmp/PluginC.dylib")))))
+
+    XCTAssertLessThan(pluginAIndex, pluginBIndex, "Order of options should be preserved")
+    XCTAssertLessThan(pluginBIndex, pluginCIndex, "Order of options should be preserved")
+  }
+
   func testClangModuleValidateOnce() throws {
     let flagTest = try Driver(args: ["swiftc", "-typecheck", "foo.swift"])
     guard flagTest.isFrontendArgSupported(.clangBuildSessionFile),