Start adding WasmKit support

Author: kabiroberai

tools/swift-plugin-server/.gitignore

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

tools/swift-plugin-server/CMakeLists.txt

@@ -57,5 +57,7 @@ if (SWIFT_BUILD_SWIFT_SYNTAX)
     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
     )
+  else()
+    # TODO: Add dependency on WasmKit
   endif()
 endif()

tools/swift-plugin-server/Package.swift

@@ -14,6 +14,7 @@ let package = Package(
   dependencies: [
     .package(path: "../../../swift-syntax"),
     .package(path: "../../lib/ASTGen"),
+    .package(path: "../../../wasmkit"),
   ],
   targets: [
     .target(
@@ -54,6 +55,8 @@ let package = Package(
         .product(name: "SwiftCompilerPluginMessageHandling", package: "swift-syntax"),
         "CSwiftPluginServer",
         "SwiftPluginServerSupport",
+        .product(name: "WasmKit", package: "WasmKit", condition: .when(platforms: [.linux, .windows])),
+        .product(name: "WASI", package: "WasmKit", condition: .when(platforms: [.linux, .windows])),
       ],
       swiftSettings: [.interoperabilityMode(.Cxx)]
     ),

tools/swift-plugin-server/Sources/swift-wasm-plugin-server/JSCWasmPlugin.swift

@@ -44,7 +44,7 @@ async (wasmData) => {
 }
 """
 
-final class JSCWasmPlugin: WasmPlugin {
+struct JSCWasmPlugin: WasmPlugin {
   private static let factory = JSContext()?.evaluateScript(js)
 
   private let handler: JSValue

tools/swift-plugin-server/Sources/swift-wasm-plugin-server/WasmKitPlugin.swift

@@ -0,0 +1,63 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#if !os(macOS)
+
+import Foundation
+import WasmKit
+import WASI
+
+typealias DefaultWasmPlugin = WasmKitPlugin
+
+struct WasmKitMacroRunner: WasmPlugin {
+    let instance: ModuleInstance
+    let runtime: Runtime
+
+    init(wasm: Data) throws {
+        let module = try parseWasm(bytes: Array(wasm))
+        let bridge = try WASIBridgeToHost()
+        runtime = Runtime(hostModules: bridge.hostModules)
+        instance = try runtime.instantiate(module: module)
+        _ = try bridge.start(instance, runtime: runtime)
+    }
+
+    func handleMessage(_ json: Data) throws -> Data {
+        let exports = instance.exports
+        guard case let .memory(memoryAddr) = exports["memory"] else { fatalError("bad memory") }
+        guard case let .function(malloc) = exports["wacro_malloc"] else { fatalError("bad wacro_malloc") }
+        guard case let .function(parse) = exports["wacro_parse"] else { fatalError("bad wacro_parse") }
+        guard case let .function(free) = exports["wacro_free"] else { fatalError("bad wacro_free") }
+
+        let inAddr = try malloc.invoke([.i32(UInt32(json.count))], runtime: runtime)[0].i32
+
+        runtime.store.withMemory(at: memoryAddr) { mem in
+            mem.data.replaceSubrange(Int(inAddr)..<(Int(inAddr) + json.count), with: json)
+        }
+
+        let outAddr = try parse.invoke([.i32(inAddr), .i32(UInt32(json.count))], runtime: runtime)[0].i32
+        let out = runtime.store.withMemory(at: memoryAddr) { mem in
+            let bytes = Array(mem.data[Int(outAddr)..<(Int(outAddr) + 4)])
+            let len =
+              (UInt32(bytes[0]) << 0)  |
+              (UInt32(bytes[1]) << 8)  |
+              (UInt32(bytes[2]) << 16) |
+              (UInt32(bytes[3]) << 24)
+            return Data(mem.data[(Int(outAddr) + 4)...].prefix(Int(len)))
+        }
+
+        _ = try free.invoke([.i32(outAddr)], runtime: runtime)
+
+        return out
+    }
+}
+
+#endif