JSCWasmPlugin: memoize, improve errors

Author: kabiroberai

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

@@ -11,11 +11,10 @@
 //===----------------------------------------------------------------------===//
 
 import JavaScriptCore
-import SwiftPluginServerSupport
 
-// returns: Promise<(ArrayBuffer) => ArrayBuffer>
+// returns: (wasm: ArrayBuffer) => Promise<(input: ArrayBuffer) => ArrayBuffer>
 private let js = """
-(async () => {
+async (wasmData) => {
   const mod = await WebAssembly.compile(wasmData);
   // stub WASI imports
   const imports = WebAssembly.Module.imports(mod)
@@ -38,39 +37,49 @@ private let js = """
     api.wacro_free(outAddr);
     return copy.buffer;
   }
-})()
+}
 """
 
 @available(macOS 10.15, *)
 final class JSCWasmPlugin: WasmPlugin {
+  private static let factory = JSContext()?.evaluateScript(js)
+
   private let handler: JSValue
 
   @MainActor init(wasm data: Data) async throws {
-    guard let context = JSContext() else { throw PluginServerError(message: "Could not create JSContext") }
-
-    let jsBuf = try JSValue(newBufferWithData: data, in: context)
-    context.globalObject.setObject(jsBuf, forKeyedSubscript: "wasmData")
-
-    guard let promise = context.evaluateScript(js) else {
-      throw PluginServerError(message: "Failed to load plugin")
+    guard let factory = Self.factory, let context = factory.context else {
+      throw JSCWasmError(message: "Failed to load plugin")
     }
 
-    if let error = context.exception {
-      throw PluginServerError(message: "Failed to load plugin: \(error)")
+    let jsBuf = try JSValue(newBufferWithData: data, in: context)
+    guard let promise = factory.call(withArguments: [jsBuf]), context.exception == nil else {
+      throw JSCWasmError(message: "Failed to load plugin", value: context.exception)
     }
 
     handler = try await promise.promiseValue
   }
 
   @MainActor func handleMessage(_ json: Data) throws -> Data {
-    let jsonJS = try JSValue(newBufferWithData: json, in: handler.context)
+    guard let context = handler.context else {
+      throw JSCWasmError(message: "Failed to invoke plugin")
+    }
+    let jsonJS = try JSValue(newBufferWithData: json, in: context)
     guard let result = handler.call(withArguments: [jsonJS]) else {
-      throw PluginServerError(message: "Wasm plugin did not provide a valid response")
+      throw JSCWasmError(message: "Wasm plugin did not provide a valid response", value: context.exception)
     }
     return result.arrayBufferData()
   }
 }
 
+public struct JSCWasmError: Error, CustomStringConvertible {
+  public let value: JSValue?
+  public let description: String
+  public init(message: String, value: JSValue? = nil) {
+    self.value = value
+    self.description = "\(message)\(value.map { ": \($0)" } ?? "")"
+  }
+}
+
 extension JSValue {
   fileprivate convenience init(newBufferWithData data: Data, in context: JSContext) throws {
     let copy = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: data.count)
@@ -105,18 +114,18 @@ extension JSValue {
         let rejectFunc = JSValue(
           object: { error in
             continuation.resume(
-              throwing: PluginServerError(message: "\(error)")
+              throwing: JSCWasmError(message: "Promise rejected", value: error)
             )
           } as @convention(block) (JSValue) -> Void,
           in: context
         )
         guard let successFunc, let rejectFunc else {
-          continuation.resume(throwing: PluginServerError(message: "Could not await promise"))
+          continuation.resume(throwing: JSCWasmError(message: "Could not await promise"))
           return
         }
         invokeMethod("then", withArguments: [successFunc, rejectFunc])
         if let exception = context.exception {
-          continuation.resume(throwing: PluginServerError(message: "\(exception)"))
+          continuation.resume(throwing: JSCWasmError(message: "Promise.then threw", value: exception))
         }
       }
     }