Add support for WebAssembly Macros (#2623)

This PR allows `CompilerPlugin`s to be built with a wasm32-unknown-wasi target, enabling them to be invoked by the Wasm Plugin runtime in swiftlang/swift#73031.

I've chosen to create a `swift_wasm_macro_pump` export to allow the caller to "drive the event loop" since issuing a `read` would by default be blocking. We need nonblocking IO because some runtimes (eg JavaScriptCore) run Wasm on the same thread as the rest of the interpreted code (JavaScript).

To test this, one can build an example macro with
```bash
Examples$ swift build \
  --experimental-swift-sdk wasm32-unknown-wasi \
  --product MacroExamplesImplementation \
  -c release
```

And then, with the changes from the swift and swift-driver PRs, a client can be compiled with
```bash
Examples$ swiftc Client.swift \
  -load-plugin-executable .build/release/MacroExamplesImplementation.wasm#MacroExamplesImplementation
```

Author: kabiroberai

Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

@@ -74,15 +74,21 @@ struct HostCapability {
 ///
 /// The low level connection and the provider is injected by the client.
 @_spi(PluginMessage)
-public class CompilerPluginMessageListener<Connection: MessageConnection, Provider: PluginProvider> {
+public class CompilerPluginMessageListener<Connection: MessageConnection, Handler: PluginMessageHandler> {
   /// Message channel for bidirectional communication with the plugin host.
   let connection: Connection
 
-  let handler: CompilerPluginMessageHandler<Provider>
+  let handler: Handler
 
-  public init(connection: Connection, provider: Provider) {
+  public init(connection: Connection, messageHandler: Handler) {
     self.connection = connection
-    self.handler = CompilerPluginMessageHandler(provider: provider)
+    self.handler = messageHandler
+  }
+
+  public init<Provider: PluginProvider>(connection: Connection, provider: Provider)
+  where Handler == PluginProviderMessageHandler<Provider> {
+    self.connection = connection
+    self.handler = PluginProviderMessageHandler(provider: provider)
   }
 
   /// Run the main message listener loop.
@@ -91,11 +97,26 @@ public class CompilerPluginMessageListener<Connection: MessageConnection, Provid
   /// On internal errors, such as I/O errors or JSON serialization errors, print
   /// an error message and `exit(1)`
   public func main() {
+    #if os(WASI)
+    // Rather than blocking on read(), let the host tell us when there's data.
+    readabilityHandler = { _ = self.handleNextMessage() }
+    #else
+    while handleNextMessage() {}
+    #endif
+  }
+
+  /// Receives and handles a single message from the plugin host.
+  ///
+  /// - Returns: `true` if there was a message to read, `false`
+  /// if the end-of-file was reached.
+  private func handleNextMessage() -> Bool {
     do {
-      while let message = try connection.waitForNextMessage(HostToPluginMessage.self) {
-        let result = handler.handleMessage(message)
-        try connection.sendMessage(result)
+      guard let message = try connection.waitForNextMessage(HostToPluginMessage.self) else {
+        return false
       }
+      let result = handler.handleMessage(message)
+      try connection.sendMessage(result)
+      return true
     } catch {
       // Emit a diagnostic and indicate failure to the plugin host,
       // and exit with an error code.
@@ -105,10 +126,18 @@ public class CompilerPluginMessageListener<Connection: MessageConnection, Provid
   }
 }
 
-/// 'CompilerPluginMessageHandler' is a type that handle a message and do the
-/// corresponding operation.
+/// A type that handles a plugin message and returns a response.
+///
+/// - SeeAlso: ``PluginProviderMessageHandler``
+@_spi(PluginMessage)
+public protocol PluginMessageHandler {
+  /// Handles a single message received from the plugin host.
+  func handleMessage(_ message: HostToPluginMessage) -> PluginToHostMessage
+}
+
+/// A `PluginMessageHandler` that uses a `PluginProvider`.
 @_spi(PluginMessage)
-public class CompilerPluginMessageHandler<Provider: PluginProvider> {
+public class PluginProviderMessageHandler<Provider: PluginProvider>: PluginMessageHandler {
   /// Object to provide actual plugin functions.
   let provider: Provider
 
@@ -199,6 +228,10 @@ public class CompilerPluginMessageHandler<Provider: PluginProvider> {
   }
 }
 
+@_spi(PluginMessage)
+@available(*, deprecated, renamed: "PluginProviderMessageHandler")
+public typealias CompilerPluginMessageHandler<Provider: PluginProvider> = PluginProviderMessageHandler<Provider>
+
 struct UnimplementedError: Error, CustomStringConvertible {
   var description: String { "unimplemented" }
 }
@@ -216,3 +249,31 @@ extension PluginProvider {
     throw UnimplementedError()
   }
 }
+
+#if compiler(>=6) && os(WASI)
+
+/// A callback invoked by the Wasm Host when new data is available on `stdin`.
+///
+/// This is safe to access without serialization as Wasm plugins are single-threaded.
+nonisolated(unsafe) private var readabilityHandler: () -> Void = {
+  fatalError(
+    """
+    CompilerPlugin.main wasn't called. Did you annotate your plugin with '@main'?
+    """
+  )
+}
+
+@_expose(wasm, "swift_wasm_macro_v1_pump")
+@_cdecl("swift_wasm_macro_v1_pump")
+func wasmPump() {
+  readabilityHandler()
+}
+
+// we can't nest the whole #if-#else in '#if os(WASI)' due to a bug where
+// '#if compiler' directives have to be the top-level #if, otherwise
+// the compiler doesn't skip unknown syntax.
+#elseif os(WASI)
+
+#error("Building swift-syntax for WebAssembly requires compiler version 6.0 or higher.")
+
+#endif

Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -26,7 +26,7 @@ import SwiftSyntax
 @_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacros
 #endif
 
-extension CompilerPluginMessageHandler {
+extension PluginProviderMessageHandler {
   /// Get concrete macro type from a pair of module name and type name.
   private func resolveMacro(_ ref: PluginMessage.MacroReference) throws -> Macro.Type {
     try provider.resolveMacro(moduleName: ref.moduleName, typeName: ref.typeName)

Sources/SwiftCompilerPluginMessageHandling/PluginMessages.swift

@@ -229,7 +229,7 @@ public enum PluginMessage {
     public var notes: [Note]
     public var fixIts: [FixIt]
 
-    internal init(
+    public init(
       message: String,
       severity: Severity,
       position: Position,

Sources/SwiftCompilerPluginMessageHandling/StandardIOMessageConnection.swift

@@ -40,6 +40,16 @@ public struct StandardIOMessageConnection: MessageConnection {
     self.outputFileDescriptor = outputFileDescriptor
   }
 
+  #if os(WASI)
+  /// Convenience initializer for Wasm executable plugins. Connects
+  /// directly to `stdin` and `stdout` as WASI doesn't support
+  /// `dup{,2}`.
+  public init() throws {
+    let inputFD = fileno(_stdin)
+    let outputFD = fileno(_stdout)
+    self.init(inputFileDescriptor: inputFD, outputFileDescriptor: outputFD)
+  }
+  #else
   /// Convenience initializer for normal executable plugins. Upon creation:
   ///   - Redirect `stdout` to `stderr` so that print statements from the plugin
   ///     are treated as plain-text output
@@ -83,6 +93,7 @@ public struct StandardIOMessageConnection: MessageConnection {
 
     self.init(inputFileDescriptor: inputFD, outputFileDescriptor: outputFD)
   }
+  #endif
 
   /// Write the buffer to the file descriptor. Throws an error on failure.
   private func _write(contentsOf buffer: UnsafeRawBufferPointer) throws {