[Macros] Add support for wasm macros

include/swift/AST/SearchPathOptions.h

@@ -468,6 +468,9 @@ class SearchPathOptions {
   /// Plugin search path options.
   std::vector<PluginSearchOption> PluginSearchOpts;
 
+  /// Path to swift-wasm-plugin-server executable.
+  std::string PluginWasmServerPath;
+
   /// Don't look in for compiler-provided modules.
   bool SkipRuntimeLibraryImportPaths = false;
 

include/swift/Option/FrontendOptions.td

@@ -277,6 +277,11 @@ def Raccess_note_EQ : Joined<["-"], "Raccess-note=">,
 def block_list_file
   : Separate<["-"], "blocklist-file">, MetaVarName<"<path>">,
     HelpText<"The path to a blocklist configuration file">;
+
+def wasm_plugin_server_path : Separate<["-"], "wasm-plugin-server-path">,
+  Flags<[DoesNotAffectIncrementalBuild, ArgumentIsPath]>,
+  HelpText<"Path to the swift-wasm-plugin-server executable">,
+  MetaVarName<"<path>">;
 } // end let Flags = [FrontendOption, NoDriverOption]
 
 def debug_crash_Group : OptionGroup<"<automatic crashing options>">;

lib/AST/PluginLoader.cpp

@@ -116,8 +116,18 @@ PluginLoader::getPluginMap() {
     case PluginSearchOption::Kind::LoadPluginExecutable: {
       auto &val = entry.get<PluginSearchOption::LoadPluginExecutable>();
       assert(!val.ExecutablePath.empty() && "empty plugin path");
-      for (auto &moduleName : val.ModuleNames) {
-        try_emplace(moduleName, /*libraryPath=*/"", val.ExecutablePath);
+      if (llvm::sys::path::filename(val.ExecutablePath).ends_with(".wasm")) {
+        // we treat wasm plugins like library plugins that can be loaded by an external
+        // "wasm server" that in turn invokes the wasm runtime.
+        const auto &wasmServerPath = Ctx.SearchPathOpts.PluginWasmServerPath;
+        assert(!wasmServerPath.empty() && "wasm load requested but got empty wasm server path");
+        for (auto &moduleName : val.ModuleNames) {
+          try_emplace(moduleName, val.ExecutablePath, wasmServerPath);
+        }
+      } else {
+        for (auto &moduleName : val.ModuleNames) {
+          try_emplace(moduleName, /*libraryPath=*/"", val.ExecutablePath);
+        }
       }
       continue;
     }

lib/Basic/Sandbox.cpp

@@ -24,8 +24,8 @@ static StringRef sandboxProfile(llvm::BumpPtrAllocator &Alloc) {
   // Allow reading file metadata of any files.
   contents += "(allow file-read-metadata)\n";
 
-  // Allow reading dylibs.
-  contents += "(allow file-read* (regex #\"\\.dylib$\"))\n";
+  // Allow reading dylibs and WebAssembly macros.
+  contents += "(allow file-read* (regex #\"\\.(dylib|wasm)$\"))\n";
 
   // This is required to launch any processes (execve(2)).
   contents += "(allow process-exec*)\n";

lib/Frontend/CompilerInvocation.cpp

@@ -1991,6 +1991,10 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts, ArgList &Args,
     }
   }
 
+  if (const Arg *A = Args.getLastArg(OPT_wasm_plugin_server_path)) {
+    Opts.PluginWasmServerPath = A->getValue();
+  }
+
   for (const Arg *A : Args.filtered(OPT_L)) {
     Opts.LibrarySearchPaths.push_back(resolveSearchPath(A->getValue()));
   }

tools/swift-plugin-server/.gitignore

@@ -0,0 +1 @@
+Package.resolved

tools/swift-plugin-server/CMakeLists.txt

@@ -1,4 +1,17 @@
 if (SWIFT_BUILD_SWIFT_SYNTAX)
+  # hack: override macOS deployment target to 10.15.4
+  # since that's the minimum required by WasmKit
+  set(SWIFT_DARWIN_DEPLOYMENT_VERSION_OSX "10.15.4")
+  include(DarwinSDKs)
+  swift_get_host_triple(SWIFT_HOST_TRIPLE)
+
+  # override the remote SwiftSystem and ArgParser dependencies in WasmKit
+  FetchContent_Declare(SwiftSystem SOURCE_DIR "${PROJECT_SOURCE_DIR}/../swift-system")
+  FetchContent_Declare(ArgumentParser SOURCE_DIR "${PROJECT_SOURCE_DIR}/../swift-argument-parser")
+
+  FetchContent_Declare(WasmKit SOURCE_DIR "${PROJECT_SOURCE_DIR}/../wasmkit")
+  FetchContent_MakeAvailable(WasmKit)
+
   # _swiftCSwiftPluginServer is just a C support library for swift-plugin-server
   # Don't bother to create '.a' for that.
   add_swift_host_library(_swiftCSwiftPluginServer OBJECT
@@ -11,6 +24,16 @@ if (SWIFT_BUILD_SWIFT_SYNTAX)
     Sources/CSwiftPluginServer/include
   )
 
+  add_pure_swift_host_library(SwiftPluginServerSupport EMIT_MODULE
+    Sources/SwiftPluginServerSupport/SwiftPluginServerSupport.swift
+    DEPENDENCIES
+      $<TARGET_OBJECTS:_swiftCSwiftPluginServer>
+      SwiftCompilerPluginMessageHandling
+  )
+  target_include_directories(SwiftPluginServerSupport PRIVATE
+    Sources/CSwiftPluginServer/include
+  )
+
   add_pure_swift_host_tool(swift-plugin-server
     Sources/swift-plugin-server/swift-plugin-server.swift
     DEPENDENCIES
@@ -23,8 +46,50 @@ if (SWIFT_BUILD_SWIFT_SYNTAX)
       SwiftSyntaxMacroExpansion
       SwiftCompilerPluginMessageHandling
       swiftLLVMJSON
+      SwiftPluginServerSupport
   )
   target_include_directories(swift-plugin-server PRIVATE
     Sources/CSwiftPluginServer/include
   )
+
+  set(SWIFT_WASM_ALLOW_JSC ON CACHE BOOL "Use JavaScriptCore Wasm runtime if possible" FORCE)
+
+  if(SWIFT_HOST_VARIANT_SDK IN_LIST SWIFT_DARWIN_PLATFORMS)
+    set(swift_wasm_use_jsc ${SWIFT_WASM_ALLOW_JSC})
+  else()
+    set(swift_wasm_use_jsc OFF)
+  endif()
+
+  if (swift_wasm_use_jsc)
+    set(wasi_dep WASI)
+    set(jsc_defines SWIFT_WASM_USE_JSC)
+  else()
+    set(wasi_dep WasmKitWASI)
+    set(jsc_defines)
+  endif()
+
+  add_pure_swift_host_tool(swift-wasm-plugin-server
+    Sources/swift-wasm-plugin-server/swift-wasm-plugin-server.swift
+    Sources/swift-wasm-plugin-server/JSCWasmEngine.swift
+    Sources/swift-wasm-plugin-server/WasmEngine.swift
+    Sources/swift-wasm-plugin-server/WasmKitEngine.swift
+    DEPENDENCIES
+      $<TARGET_OBJECTS:_swiftCSwiftPluginServer>
+    SWIFT_COMPONENT
+      compiler
+    SWIFT_DEPENDENCIES
+      SwiftCompilerPluginMessageHandling
+      swiftLLVMJSON
+      SwiftPluginServerSupport
+      ${wasi_dep}
+  )
+  target_compile_definitions(swift-wasm-plugin-server PRIVATE ${jsc_defines})
+  target_include_directories(swift-wasm-plugin-server PRIVATE
+    Sources/CSwiftPluginServer/include
+  )
+  if (swift_wasm_use_jsc)
+    add_custom_command(TARGET swift-wasm-plugin-server POST_BUILD COMMAND
+      codesign -fs - $<TARGET_FILE:swift-wasm-plugin-server> --entitlements ${CMAKE_CURRENT_SOURCE_DIR}/allow-jit.plist
+    )
+  endif()
 endif()

tools/swift-plugin-server/Package.swift

@@ -1,38 +1,68 @@
-// swift-tools-version: 5.6
+// swift-tools-version: 5.9
 
 import PackageDescription
 
+let allowJSC = true
+
 let package = Package(
   name: "swift-plugin-server",
   platforms: [
-    .macOS(.v10_15)
+    .macOS(.v11),
+  ],
+  products: [
+    .library(name: "swift-plugin-server", targets: ["swift-plugin-server"]),
+    .library(name: "swift-wasm-plugin-server", targets: ["swift-wasm-plugin-server"]),
   ],
   dependencies: [
     .package(path: "../../../swift-syntax"),
     .package(path: "../../lib/ASTGen"),
+    .package(path: "../../../wasmkit"),
   ],
   targets: [
     .target(
       name: "CSwiftPluginServer",
       cxxSettings: [
         .unsafeFlags([
           "-I", "../../include",
+          "-I", "../../stdlib/public/SwiftShims",
+          "-I", "../../../build/Default/swift/include",
+          "-I", "../../../build/Default/llvm/include",
           "-I", "../../../llvm-project/llvm/include",
         ])
       ]
     ),
-    .executableTarget(
-      name: "swift-plugin-server",
+    .target(
+      name: "SwiftPluginServerSupport",
       dependencies: [
         .product(name: "swiftLLVMJSON", package: "ASTGen"),
         .product(name: "SwiftCompilerPluginMessageHandling", package: "swift-syntax"),
-        .product(name: "SwiftDiagnostics", package: "swift-syntax"),
-        .product(name: "SwiftSyntax", package: "swift-syntax"),
-        .product(name: "SwiftOperators", package: "swift-syntax"),
-        .product(name: "SwiftParser", package: "swift-syntax"),
+        "CSwiftPluginServer",
+      ],
+      swiftSettings: [.interoperabilityMode(.Cxx)]
+    ),
+    .target(
+      name: "swift-plugin-server",
+      dependencies: [
+        .product(name: "SwiftCompilerPluginMessageHandling", package: "swift-syntax"),
         .product(name: "SwiftSyntaxMacros", package: "swift-syntax"),
-        "CSwiftPluginServer"
-      ]
+        "CSwiftPluginServer",
+        "SwiftPluginServerSupport",
+      ],
+      swiftSettings: [.interoperabilityMode(.Cxx)]
+    ),
+    .target(
+      name: "swift-wasm-plugin-server",
+      dependencies: [
+        .product(name: "swiftLLVMJSON", package: "ASTGen"),
+        .product(name: "SwiftCompilerPluginMessageHandling", package: "swift-syntax"),
+        "CSwiftPluginServer",
+        "SwiftPluginServerSupport",
+        .product(name: "WASI", package: "WasmKit"),
+        .product(name: "WasmKitWASI", package: "WasmKit", condition: .when(platforms: [.linux, .windows])),
+      ],
+      swiftSettings: [.interoperabilityMode(.Cxx)] + (
+        allowJSC ? [.define("SWIFT_WASM_USE_JSC", .when(platforms: [.macOS]))] : []
+      )
     ),
   ],
   cxxLanguageStandard: .cxx17

tools/swift-plugin-server/Sources/SwiftPluginServerSupport/SwiftPluginServerSupport.swift

@@ -0,0 +1,137 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2024 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
+//
+//===----------------------------------------------------------------------===//
+
+@_spi(PluginMessage) import SwiftCompilerPluginMessageHandling
+import swiftLLVMJSON
+import CSwiftPluginServer
+
+@_spi(PluginMessage) public final class PluginHostConnection: MessageConnection {
+  let handle: UnsafeRawPointer
+  public init() throws {
+    var errorMessage: UnsafePointer<CChar>? = nil
+    guard let handle = PluginServer_createConnection(&errorMessage) else {
+      throw PluginServerError(message: String(cString: errorMessage!))
+    }
+    self.handle = handle
+  }
+
+  deinit {
+    PluginServer_destroyConnection(self.handle)
+  }
+
+  public func sendMessage<TX: Encodable>(_ message: TX) throws {
+    try LLVMJSON.encoding(message) { buffer in
+      try self.sendMessageData(buffer)
+    }
+  }
+
+  public func waitForNextMessage<RX: Decodable>(_ type: RX.Type) throws -> RX? {
+    return try self.withReadingMessageData { jsonData in
+      try LLVMJSON.decode(RX.self, from: jsonData)
+    }
+  }
+
+  /// Send a serialized message to the message channel.
+  private func sendMessageData(_ data: UnsafeBufferPointer<Int8>) throws {
+    // Write the header (a 64-bit length field in little endian byte order).
+    var header: UInt64 = UInt64(data.count).littleEndian
+    let writtenSize = try Swift.withUnsafeBytes(of: &header) { buffer in
+      try self.write(buffer: UnsafeRawBufferPointer(buffer))
+    }
+    guard writtenSize == MemoryLayout.size(ofValue: header) else {
+      throw PluginServerError(message: "failed to write message header")
+    }
+
+    // Write the body.
+    guard try self.write(buffer: UnsafeRawBufferPointer(data)) == data.count else {
+      throw PluginServerError(message: "failed to write message body")
+    }
+  }
+
+  /// Read a serialized message from the message channel and call the 'body'
+  /// with the data.
+  private func withReadingMessageData<R>(_ body: (UnsafeBufferPointer<Int8>) throws -> R) throws -> R? {
+    // Read the header (a 64-bit length field in little endian byte order).
+    var header: UInt64 = 0
+    let readSize = try Swift.withUnsafeMutableBytes(of: &header) { buffer in
+      try self.read(into: UnsafeMutableRawBufferPointer(buffer))
+    }
+    guard readSize == MemoryLayout.size(ofValue: header) else {
+      if readSize == 0 {
+        // The host closed the pipe.
+        return nil
+      }
+      // Otherwise, some error happened.
+      throw PluginServerError(message: "failed to read message header")
+    }
+
+    // Read the body.
+    let count = Int(UInt64(littleEndian: header))
+    let data = UnsafeMutableBufferPointer<Int8>.allocate(capacity: count)
+    defer { data.deallocate() }
+    guard try self.read(into: UnsafeMutableRawBufferPointer(data)) == count else {
+      throw PluginServerError(message: "failed to read message body")
+    }
+
+    // Invoke the handler.
+    return try body(UnsafeBufferPointer(data))
+  }
+
+  /// Write the 'buffer' to the message channel.
+  /// Returns the number of bytes succeeded to write.
+  private func write(buffer: UnsafeRawBufferPointer) throws -> Int {
+    var bytesToWrite = buffer.count
+    guard bytesToWrite > 0 else {
+      return 0
+    }
+    var ptr = buffer.baseAddress!
+
+    while (bytesToWrite > 0) {
+      let writtenSize = PluginServer_write(handle, ptr, bytesToWrite)
+      if (writtenSize <= 0) {
+        // error e.g. broken pipe.
+        break
+      }
+      ptr = ptr.advanced(by: writtenSize)
+      bytesToWrite -= Int(writtenSize)
+    }
+    return buffer.count - bytesToWrite
+  }
+
+  /// Read data from the message channel into the 'buffer' up to 'buffer.count' bytes.
+  /// Returns the number of bytes succeeded to read.
+  private func read(into buffer: UnsafeMutableRawBufferPointer) throws -> Int {
+    var bytesToRead = buffer.count
+    guard bytesToRead > 0 else {
+      return 0
+    }
+    var ptr = buffer.baseAddress!
+
+    while bytesToRead > 0 {
+      let readSize = PluginServer_read(handle, ptr, bytesToRead)
+      if (readSize <= 0) {
+        // 0: EOF (the host closed), -1: Broken pipe (the host crashed?)
+        break;
+      }
+      ptr = ptr.advanced(by: readSize)
+      bytesToRead -= readSize
+    }
+    return buffer.count - bytesToRead
+  }
+}
+
+public struct PluginServerError: Error, CustomStringConvertible {
+  public var description: String
+  public init(message: String) {
+    self.description = message
+  }
+}