Add support for WebAssembly Macros #2623

kabiroberai · 2024-04-24T19:54:25Z

This PR allows CompilerPlugins to be built with a wasm32-unknown-wasi target, enabling them to be invoked by the Wasm Plugin runtime in apple/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

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

Examples$ swiftc Client.swift \
  -load-plugin-executable .build/release/MacroExamplesImplementation.wasm#MacroExamplesImplementation

ahoppen

Looks good from my side but I would like to get a review from @rintaro as well.

ahoppen · 2024-04-24T21:25:01Z

Sources/SwiftCompilerPlugin/CompilerPlugin.swift

+#if compiler(>=6.0)
+
+@_expose(wasm, "swift_wasm_macro_pump")
+@_cdecl("swift_wasm_macro_pump")
+func wasmPump() {
+  readabilityHandler()
+}
+
+#endif


Does this code require Swift 6 to build for wasm? What happens if you build using a Swift 5 compiler? If it doesn’t build with an older compiler, I think we should put a #error in the #else branch telling you that you need a Swift 6 compiler to build the plugin server for WASM.

Turns out this was causing pre-6.0 builds (even those not targeting wasm) to fail — somehow Swift ignores the #if compiler check and complains that it doesn't recognize @_expose(wasm). I moved the export over to C, so this should now work fine with older compilers. That said I haven't tested a pre-6.0 Wasm-compatible compiler yet, since afaik Swift didn't officially support Wasm until 6.0 — but I can give it a spin with a SwiftWasm toolchain.

kabiroberai · 2024-04-24T22:20:51Z

A couple of things worth mentioning:

I didn't add wasm_support.c to the CMake source set for _SwiftSyntaxCShims since I don't think we ever build SwiftSyntax for WASI with CMake — and I'm not sure how I would even test that. Is this okay or should I try to allow cross-compiling with CMake?
Can we add a smoke test to ensure that SwiftSyntax continues to build for WASI? Unsure of the best way to go about doing this.

kateinoigakukun · 2024-04-25T04:35:44Z

@kabiroberai

Can we add a smoke test to ensure that SwiftSyntax continues to build for WASI? Unsure of the best way to go about doing this.

I'm not sure we have a way to run smoke test with Wasm target in CI right now. But we are trying to build Wasm SDK in CI here apple/swift#72728 and after the integration, we will be able to use Wasm target in CI.

rintaro · 2024-04-29T04:41:09Z

Sources/SwiftCompilerPluginMessageHandling/StandardIOMessageConnection.swift

+    // Wasm doesn't support dup{,2} so we use the original file descriptors.
+    let inputFD = fileno(_stdin)
+    let outputFD = fileno(_stdout)
+    #else


Instead of #if in the body, could you just wrap the whole decl? The body for wasm would be just

#if os(WASI) public init() throws { self.init( inputFileDescriptor: fileno(_stdin), outputFileDescriptor: fileno(_stdout) ) } #else ...

This way we can find #endif much easier.

rintaro · 2024-04-29T04:59:03Z

Sources/_SwiftSyntaxCShims/wasm_support.c

+// 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");


Could you teach me what this does exactly?

This adds a custom swift_wasm_macro_abi section to the binary with a little-endian uint32 value. We check the ABI version in the wasm executor in order to future-proof. See:

https://github.com/apple/swift/pull/73031/files#diff-2261ec558c289a4bd568f211772e906caa997f8c26b38e4c6d5d37b7b1554379R47

Will add a brief explanation to the comment.

hmm actually @kateinoigakukun what do you think about using the export name to indicate the ABI instead? E.g. we could export the pump function with the name swift_wasm_macro_pump_v1 and the runtime could check for the presence of this to indicate the v1 ABI. The custom section support feels shaky to me — I wouldn't be surprised if there are tools that don't know how to handle custom wasm sections (or, alternatively, strip them out.)

Pushed this as a tentative change, happy to revert if the Custom Sections approach is preferable but I do like this atm because it's a bit more lightweight.

For posterity: discussed this with @kateinoigakukun here; we aligned on using the exported function name rather than a custom section, since some Wasm tooling (ahem wasm-strip) can be too trigger-happy about stripping custom sections.

I think Swift tries to be conservative with custom sections because it's a terribly non-portable solution. Although portability isn't at issue here, it might maybe make sense to use an exported function name. That'd be the typical solution here if we were talking about a C function on another platform, I think.

rintaro · 2024-04-29T05:00:55Z

Sources/SwiftCompilerPluginMessageHandling/StandardIOMessageConnection.swift

+// fatalError writes to stdout, which is the message
+// output stream under WASI
+public func internalError(_ message: String) -> Never {
+  fputs("Internal Error: \(message)\n", _stderr)


fatalError writes to stdout

As far as I can see it's stderr. Am I missing anything? https://github.com/apple/swift/blob/77f53a5e50f0cf964ae9dcc42f78ec43227a3db2/stdlib/public/runtime/Errors.cpp#L323-L332

But anyway, I think I'm going to use fputs("Internal Error: \(message)\n", _stderr) in all platforms, (we don't want #file etc, and other magic in fatalError). So I will bring internalError() back to CompilerPlugin.swift

Indeed, looks like you're right. Not sure why I thought it wrote to stdout — might've forgotten to attach stderr to the WASI bridge at some point.

rintaro · 2024-04-29T05:07:08Z

Sources/SwiftCompilerPlugin/CompilerPlugin.swift

+    #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


I feel this logic should be sinked to impl.main(), so handleNextMessage() can be internal.

hmm I'm worried that moving readabilityHandler into SwiftCompilerPluginMessageHandling would be leaky, since the swift_wasm_macro_pump export (which invokes readabilityHandler) is an impl detail of the SwiftCompilerPlugin target. Also worth mentioning that CompilerPluginMessageListener is @_spi(PluginMessage) so I don't think we're creating any new API contracts even if handleNextMessage is technically public.

Actually I just realized that _SwiftSyntaxCShims is a dependency of SwiftCompilerPluginMessageHandling so the latter is already strongly coupled to our wasm ABI. Moved readabilityHandler into that module.

rintaro · 2024-04-29T05:39:10Z

Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

-    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


I opened a PR #2631

My intention is, in this PR, you'd make it like:

public func main() { #if os(WASI) readabilityHandler = { _ = self.handleNextMessage() } #else while handleNextMessage() {} #endif } func handleNextMessage() -> Bool { do { guard let message = try connection.waitForNextMessage(HostToPluginMessage.self) else { return false } let result = handler.handleMessage(message) try connection.sendMessage(result) return true } catch { fputs("Internal Error: \(message)\n", _stderr) exit(1) } }

I think this should minimize #if os(WASI) branch code and improves the readability.

# Conflicts: # Sources/SwiftCompilerPlugin/CompilerPlugin.swift # Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

kabiroberai · 2024-05-31T00:41:14Z

Hey @rintaro, could you please take another look at this when you have a minute? Thanks!

kabiroberai requested review from ahoppen and bnbarham as code owners April 24, 2024 19:54

kabiroberai marked this pull request as draft April 24, 2024 20:22

kabiroberai marked this pull request as ready for review April 24, 2024 21:12

ahoppen reviewed Apr 24, 2024

View reviewed changes

ahoppen requested a review from rintaro April 24, 2024 21:26

Rebase

158ddd5

kabiroberai force-pushed the kabir/wasm-plugins branch from 9305330 to 158ddd5 Compare April 28, 2024 21:22

kabiroberai mentioned this pull request Apr 28, 2024

[Macros] Add support for wasm macros apple/swift#73031

Draft

4 tasks

rintaro reviewed Apr 29, 2024

View reviewed changes

kabiroberai added 8 commits April 29, 2024 23:14

Merge branch 'main' into kabir/wasm-plugins

385cc0a

# Conflicts: # Sources/SwiftCompilerPlugin/CompilerPlugin.swift # Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

cleanup

4dd83b3

We don’t need internalError

0b6a5a6

Comment

7a6b576

Feedback

f5d82ef

More feedback

9f455cb

comment

18814f1

comment

a38baa4

kabiroberai requested a review from rintaro April 30, 2024 14:39

kabiroberai added 6 commits May 5, 2024 13:56

Merge branch 'main' into kabir/wasm-plugins

a582c2f

Change wasm versioning approach

85e0374

Merge branch 'main' into kabir/wasm-plugins

921b7ea

Merge branch 'main' into kabir/wasm-plugins

9fab7cc

Create PluginMessageHandler

d3495f2

Make PluginMessage.Diagnostic.init SPI

76de84f

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for WebAssembly Macros #2623

Add support for WebAssembly Macros #2623

kabiroberai commented Apr 24, 2024 •

edited

ahoppen left a comment

ahoppen Apr 24, 2024

kabiroberai Apr 24, 2024

kabiroberai commented Apr 24, 2024

kateinoigakukun commented Apr 25, 2024

rintaro Apr 29, 2024 •

edited

rintaro Apr 29, 2024

kabiroberai Apr 30, 2024

kabiroberai May 3, 2024

kabiroberai May 5, 2024

kabiroberai May 11, 2024

grynspan May 16, 2024

rintaro Apr 29, 2024

kabiroberai Apr 30, 2024

rintaro Apr 29, 2024

kabiroberai Apr 30, 2024

kabiroberai Apr 30, 2024

rintaro Apr 29, 2024 •

edited

kabiroberai commented May 31, 2024

Add support for WebAssembly Macros #2623

Are you sure you want to change the base?

Add support for WebAssembly Macros #2623

Conversation

kabiroberai commented Apr 24, 2024 • edited

ahoppen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kabiroberai commented Apr 24, 2024

kateinoigakukun commented Apr 25, 2024

rintaro Apr 29, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rintaro Apr 29, 2024 • edited

Choose a reason for hiding this comment

kabiroberai commented May 31, 2024

kabiroberai commented Apr 24, 2024 •

edited

rintaro Apr 29, 2024 •

edited

rintaro Apr 29, 2024 •

edited