Add support for WebAssembly Macros

Sources/SwiftCompilerPlugin/CompilerPlugin.swift

@@ -112,12 +112,45 @@ extension CompilerPlugin {
     let connection = try StandardIOMessageConnection()
     let provider = MacroProviderAdapter(plugin: Self())
     let impl = CompilerPluginMessageListener(connection: connection, provider: provider)
+    #if os(WASI)
+    // Rather than blocking on read(), let the host tell us when there's data.
+    readabilityHandler = {
+      do {
+        _ = try impl.handleNextMessage()
+      } catch {
+        internalError("\(error)")
+      }
+    }
+    #else
     do {
       try impl.main()
     } catch {
       // Emit a diagnostic and indicate failure to the plugin host,
       // and exit with an error code.
       fatalError("Internal Error: \(error)")
     }
+    #endif
   }
 }
+
+#if 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 = {
+  internalError("""
+  CompilerPlugin.main wasn't called. Did you annotate your plugin with '@main'?
+  """)
+}
+
+// We can use @_expose(wasm, ...) with compiler >= 6.0
+// but it causes build errors with older compilers.
+// Instead, we export from wasm_support.c and trampoline
+// to this method.
+@_cdecl("_swift_wasm_macro_pump")
+func wasmPump() {
+  readabilityHandler()
+}
+
+#endif

Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

@@ -86,10 +86,20 @@ public class CompilerPluginMessageListener<Connection: MessageConnection, Provid
   /// Throws an error when it failed to send/receive the message, or failed
   /// to serialize/deserialize the message.
   public func main() throws {
-    while let message = try connection.waitForNextMessage(HostToPluginMessage.self) {
-      let result = handler.handleMessage(message)
-      try connection.sendMessage(result)
+    while try handleNextMessage() {}
+  }
+
+  /// 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.
+  public func handleNextMessage() throws -> Bool {
+    guard let message = try connection.waitForNextMessage(HostToPluginMessage.self) else {
+      return false
     }
+    let result = handler.handleMessage(message)
+    try connection.sendMessage(result)
+    return true
   }
 }
 

Sources/SwiftCompilerPluginMessageHandling/StandardIOMessageConnection.swift

@@ -48,6 +48,11 @@ public struct StandardIOMessageConnection: MessageConnection {
   ///   - Duplicate the original `stdin` and `stdout` for use as messaging
   ///     pipes, and are not directly used by the plugin logic
   public init() throws {
+    #if os(WASI)
+    // Wasm doesn't support dup{,2} so we use the original file descriptors.
+    let inputFD = fileno(_stdin)
+    let outputFD = fileno(_stdout)
+    #else
     // Duplicate the `stdin` file descriptor, which we will then use for
     // receiving messages from the plugin host.
     let inputFD = dup(fileno(_stdin))
@@ -80,6 +85,7 @@ public struct StandardIOMessageConnection: MessageConnection {
     _ = _setmode(inputFD, _O_BINARY)
     _ = _setmode(outputFD, _O_BINARY)
     #endif
+    #endif
 
     self.init(inputFileDescriptor: inputFD, outputFileDescriptor: outputFD)
   }
@@ -146,6 +152,15 @@ public struct StandardIOMessageConnection: MessageConnection {
   }
 }
 
+#if os(WASI)
+// fatalError writes to stdout, which is the message
+// output stream under WASI
+public func internalError(_ message: String) -> Never {
+  fputs("Internal Error: \(message)\n", _stderr)
+  exit(1)
+}
+#endif
+
 private enum IOError: Error, CustomStringConvertible {
   case readReachedEndOfInput
   case systemError(function: String, errno: CInt)

Sources/_SwiftSyntaxCShims/wasm_support.c

@@ -0,0 +1,35 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#if defined(__wasm32__)
+
+// uint32_t
+#define SWIFT_WASM_MACRO_ABI 1
+
+#define _STR(X) #X
+#define STR(X) _STR(X)
+
+// LLVM has special-cased handling to map .custom_section.foo to
+// Wasm Custom Section "foo". this must be a metadata section rather
+// than a data section so we can't use __attribute__((section)) for it.
+// See: https://reviews.llvm.org/D43097
+__asm__("\t.section .custom_section.swift_wasm_macro_abi,\"\",@\n\t.4byte " STR(SWIFT_WASM_MACRO_ABI) "\n");
+
+// defined in CompilerPlugin.swift
+void _swift_wasm_macro_pump(void);
+
+__attribute__((export_name("swift_wasm_macro_pump")))
+void swift_wasm_macro_pump(void) {
+    _swift_wasm_macro_pump();
+}
+
+#endif