Change wasm versioning approach

Author: kabiroberai

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

@@ -27,7 +27,6 @@ async (wasmData, imports) => {
   const api = instance.exports;
   return {
     api,
-    customSections: (name) => WebAssembly.Module.customSections(mod, name),
     read: (off, size) => new Uint8Array(new Uint8Array(api.memory.buffer, off, size)).buffer,
     write: (buf, off, size) => new Uint8Array(api.memory.buffer, off, size).set(new Uint8Array(buf)),
   };
@@ -79,20 +78,18 @@ final class JSCWasmEngine: WasmEngine {
     memory = JSCGuestMemory(getMemory: getMemory, setMemory: setMemory)
   }
 
-  func customSections(named name: String) throws -> [ArraySlice<UInt8>] {
-    guard let array = runner.invokeMethod("customSections", withArguments: [name]),
-          let length = array.objectForKeyedSubscript("length")?.toUInt32() else { return [] }
-    return (0..<length)
-      .compactMap { array.objectAtIndexedSubscript(Int($0)) }
-      .map { $0.withUnsafeArrayBuffer { Array($0)[...] } }
-  }
-
-  func invoke(_ method: String, _ args: [UInt32]) throws -> [UInt32] {
-    let result = api.invokeMethod(method, withArguments: args)!
-    if let exception = api.context.exception {
-      throw JSCWasmError(message: "Call to \(method) failed", value: exception)
+  func function(named name: String) throws -> WasmFunction? {
+    guard let export = api.objectForKeyedSubscript(name), export.isObject else {
+      return nil
+    }
+    return { [api] args in
+      let result = export.call(withArguments: args)
+      if let exception = api.context.exception {
+        throw JSCWasmError(message: "Call to '\(name)' failed", value: exception)
+      }
+      guard let result else { return [] }
+      return result.isUndefined ? [] : [result.toUInt32()]
     }
-    return result.isUndefined ? [] : [result.toUInt32()]
   }
 }
 

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

@@ -14,12 +14,12 @@ import WASI
 import WasmTypes
 import SystemPackage
 
+typealias WasmFunction = ([UInt32]) throws -> [UInt32]
+
 protocol WasmEngine {
   init(path: FilePath, imports: WASIBridgeToHost) async throws
 
-  func customSections(named name: String) throws -> [ArraySlice<UInt8>]
-
-  func invoke(_ method: String, _ args: [UInt32]) throws -> [UInt32]
+  func function(named name: String) throws -> WasmFunction?
 }
 
 typealias DefaultWasmPlugin = WasmEnginePlugin<DefaultWasmEngine>
@@ -28,6 +28,7 @@ typealias DefaultWasmPlugin = WasmEnginePlugin<DefaultWasmEngine>
 struct WasmEnginePlugin<Engine: WasmEngine>: WasmPlugin {
   private let hostToPlugin: FileDescriptor
   private let pluginToHost: FileDescriptor
+  private let pumpFunction: WasmFunction
   let engine: Engine
 
   init(path: FilePath) async throws {
@@ -42,39 +43,17 @@ struct WasmEnginePlugin<Engine: WasmEngine>: WasmPlugin {
       stderr: .standardError
     )
     engine = try await Engine(path: path, imports: bridge)
-    try checkABIVersion()
-    _ = try engine.invoke("_start", [])
-  }
 
-  private func checkABIVersion() throws {
-    let abiVersion = try abiVersion()
-    guard abiVersion == 1 else {
-      throw WasmEngineError(message: "Wasm plugin has unsupported ABI version: \(abiVersion)")
+    let exportName = "swift_wasm_macro_v1_pump"
+    guard let pump = try engine.function(named: exportName) else {
+      throw WasmEngineError(message: "Wasm plugin has an unknown ABI (could not find '\(exportName)')")
     }
-  }
+    self.pumpFunction = pump
 
-  private func abiVersion() throws -> UInt32 {
-    let sectionName = "swift_wasm_macro_abi"
-    let sections = try engine.customSections(named: sectionName)
-    switch sections.count {
-    case 0:
-      throw WasmEngineError(message: "Wasm macro is missing a '\(sectionName)' section")
-    case 1:
-      break
-    default:
-      throw WasmEngineError(message: "Wasm macro has too many '\(sectionName)' sections. Expected one, got \(sections.count)")
-    }
-    let section = sections[0]
-    guard section.count == 4 else {
-      throw WasmEngineError(message: """
-      Wasm macro has incorrect '\(sectionName)' section length. Expected 4 bytes, got \(section.count).
-      """)
-    }
-    return section.withUnsafeBufferPointer { buffer in
-      buffer.withMemoryRebound(to: UInt32.self) {
-        UInt32(littleEndian: $0.baseAddress!.pointee)
-      }
+    guard let start = try engine.function(named: "_start") else {
+      throw WasmEngineError(message: "Wasm plugin does not have a '_start' entrypoint")
     }
+    _ = try start([])
   }
 
   func handleMessage(_ json: [UInt8]) async throws -> [UInt8] {
@@ -83,7 +62,7 @@ struct WasmEnginePlugin<Engine: WasmEngine>: WasmPlugin {
     }
     try hostToPlugin.writeAll(json)
 
-    _ = try engine.invoke("swift_wasm_macro_pump", [])
+    _ = try pumpFunction([])
 
     let lengthRaw = try withUnsafeTemporaryAllocation(of: UInt8.self, capacity: 8) { buffer in
       let lengthCount = try pluginToHost.read(into: UnsafeMutableRawBufferPointer(buffer))

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

@@ -32,12 +32,11 @@ struct WasmKitEngine: WasmEngine {
     instance = try runtime.instantiate(module: module)
   }
 
-  func customSections(named name: String) throws -> [ArraySlice<UInt8>] {
-    module.customSections.filter { $0.name == name }.map(\.bytes)
-  }
-
-  func invoke(_ method: String, _ args: [UInt32]) throws -> [UInt32] {
-    try runtime.invoke(instance, function: method, with: args.map(Value.i32)).map(\.i32)
+  func function(named name: String) throws -> WasmFunction? {
+    guard case let .function(function) = instance.exportInstances.first(where: { $0.name == name })?.value else {
+      return nil
+    }
+    return { args in try function.invoke(args.map(Value.i32), runtime: runtime).map(\.i32) }
   }
 }