Skip to content

Instantly share code, notes, and snippets.

@andrewmd5

andrewmd5/Asyncify.swift

Created November 22, 2025 02:14
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save andrewmd5/c429a35796e1f88d5d762c2848f38e61 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save andrewmd5/c429a35796e1f88d5d762c2848f38e61 to your computer and use it in GitHub Desktop.
Download ZIP
Asyncify for WasmKit
Raw
Asyncify.swift
import Foundation
import WasmKit
/// Asyncify state machine states
private enum AsyncifyState: UInt32 {
case none = 0
case unwinding = 1
case rewinding = 2
}
/// Memory layout configuration for Asyncify data structure
private final class AsyncifyMemoryLayout: @unchecked Sendable {
private let lock = NSLock()
private var _dataAddr: UInt32 = 0
private var _dataStart: UInt32 = 0
private var _dataEnd: UInt32 = 0
private var _initialized: Bool = false
var dataAddr: UInt32 {
lock.lock()
defer { lock.unlock() }
return _dataAddr
}
var dataStart: UInt32 {
lock.lock()
defer { lock.unlock() }
return _dataStart
}
var dataEnd: UInt32 {
lock.lock()
defer { lock.unlock() }
return _dataEnd
}
func initialize(instance: Instance) throws {
lock.lock()
defer { lock.unlock() }
guard !_initialized else { return }
// Get the stack pointer to find safe memory region
if let stackPointerGlobal = instance.exports[global: "__stack_pointer"] {
let stackPointerValue = stackPointerGlobal.value
if case .i32(let stackPointer) = stackPointerValue {
let stackBase = UInt32(bitPattern: Int32(stackPointer))
_dataAddr = stackBase - 65536
_dataStart = _dataAddr + 8
_dataEnd = _dataStart + 8192
_initialized = true
return
}
}
_dataAddr = 16
_dataStart = 24
_dataEnd = 1024
_initialized = true
}
}
/// Error types specific to Asyncify operations
public enum AsyncifyError: Error, CustomStringConvertible {
case invalidState(current: UInt32, expected: UInt32)
case exportsNotInitialized
case memoryNotFound
case missingAsyncifyExport(String)
case invalidReturnValue
case instanceDeallocated
case maxIterationsExceeded
public var description: String {
switch self {
case .invalidState(let current, let expected):
return "Invalid Asyncify state: \(current), expected: \(expected)"
case .exportsNotInitialized:
return "Asyncify exports not initialized"
case .memoryNotFound:
return "Memory not found in exports or imports"
case .missingAsyncifyExport(let name):
return "Missing required Asyncify export: \(name)"
case .invalidReturnValue:
return "Invalid return value from async import"
case .instanceDeallocated:
return "Instance was deallocated during async operation"
case .maxIterationsExceeded:
return "Maximum asyncify iterations exceeded (possible infinite loop)"
}
}
}
/// Result type that can be either synchronous or asynchronous
public enum ImportResult {
case sync([Value])
case async(Task<[Value], Error>)
}
/// Thread-safe storage for the async task and its result
private final class TaskStorage: @unchecked Sendable {
private let lock = NSLock()
private var _value: Any?
func setValue(_ value: Any?) {
lock.lock()
defer { lock.unlock() }
_value = value
}
func getValue() -> Any? {
lock.lock()
defer { lock.unlock() }
return _value
}
}
/// Core Asyncify implementation for WasmKit
public final class Asyncify: @unchecked Sendable {
fileprivate let taskStorage = TaskStorage()
fileprivate let memoryLayout = AsyncifyMemoryLayout()
fileprivate var getStateFunc: Function?
fileprivate var startUnwindFunc: Function?
fileprivate var stopUnwindFunc: Function?
fileprivate var startRewindFunc: Function?
fileprivate var stopRewindFunc: Function?
fileprivate var instance: Instance?
public init() {}
fileprivate func getState() throws -> AsyncifyState {
guard let getStateFunc = getStateFunc else {
throw AsyncifyError.exportsNotInitialized
}
let result = try getStateFunc([])
guard result.count == 1, case .i32(let stateValue) = result[0] else {
throw AsyncifyError.invalidReturnValue
}
guard let state = AsyncifyState(rawValue: UInt32(bitPattern: Int32(stateValue))) else {
throw AsyncifyError.invalidState(
current: UInt32(bitPattern: Int32(stateValue)), expected: 0)
}
return state
}
fileprivate func assertNoneState() throws {
let state = try getState()
if state != .none {
throw AsyncifyError.invalidState(
current: state.rawValue, expected: AsyncifyState.none.rawValue)
}
}
/// Initialize Asyncify with the instantiated module
public func initialize(instance: Instance) throws {
self.instance = instance
try memoryLayout.initialize(instance: instance)
guard let getState = instance.exports[function: "asyncify_get_state"] else {
throw AsyncifyError.missingAsyncifyExport("asyncify_get_state")
}
guard let startUnwind = instance.exports[function: "asyncify_start_unwind"] else {
throw AsyncifyError.missingAsyncifyExport("asyncify_start_unwind")
}
guard let stopUnwind = instance.exports[function: "asyncify_stop_unwind"] else {
throw AsyncifyError.missingAsyncifyExport("asyncify_stop_unwind")
}
guard let startRewind = instance.exports[function: "asyncify_start_rewind"] else {
throw AsyncifyError.missingAsyncifyExport("asyncify_start_rewind")
}
guard let stopRewind = instance.exports[function: "asyncify_stop_rewind"] else {
throw AsyncifyError.missingAsyncifyExport("asyncify_stop_rewind")
}
guard let memory = instance.exports[memory: "memory"] else {
throw AsyncifyError.memoryNotFound
}
self.getStateFunc = getState
self.startUnwindFunc = startUnwind
self.stopUnwindFunc = stopUnwind
self.startRewindFunc = startRewind
self.stopRewindFunc = stopRewind
memory.withUnsafeMutableBufferPointer(offset: UInt(memoryLayout.dataAddr), count: 8) {
buffer in
let dataPtr = buffer.baseAddress!
dataPtr.storeBytes(of: memoryLayout.dataStart.littleEndian, as: UInt32.self)
dataPtr.advanced(by: 4).storeBytes(
of: memoryLayout.dataEnd.littleEndian, as: UInt32.self)
}
}
/// Wraps an import function that returns ImportResult to support Asyncify suspension
public func wrapImport(
_ function: @escaping @Sendable (Caller, [Value]) throws -> ImportResult
) -> @Sendable (Caller, [Value]) throws -> [Value] {
return { [weak self] caller, args in
guard let self = self else {
throw AsyncifyError.instanceDeallocated
}
let currentState = try self.getState()
if currentState == .rewinding {
guard let stopRewind = self.stopRewindFunc else {
throw AsyncifyError.exportsNotInitialized
}
_ = try stopRewind([])
if let value = self.taskStorage.getValue() {
return value as! [Value]
}
}
try self.assertNoneState()
let importResult = try function(caller, args)
switch importResult {
case .sync(let value):
return value
case .async:
guard let startUnwind = self.startUnwindFunc else {
throw AsyncifyError.exportsNotInitialized
}
self.taskStorage.setValue(importResult)
_ = try startUnwind([.i32(UInt32(Int32(bitPattern: self.memoryLayout.dataAddr)))])
return []
}
}
}
/// Wraps an export function to handle async suspension and resumption
public func wrapExport(_ function: Function) -> AsyncFunction {
AsyncFunction(asyncify: self, inner: function)
}
/// Wraps all exports in an instance to support async operations
public func wrapExports(_ instance: Instance) -> AsyncExports {
AsyncExports(asyncify: self, instance: instance)
}
}
/// An async-enabled wrapper around a WebAssembly function
public final class AsyncFunction: @unchecked Sendable {
private let asyncify: Asyncify
private let inner: Function
fileprivate init(asyncify: Asyncify, inner: Function) {
self.asyncify = asyncify
self.inner = inner
}
public func callAsFunction(_ arguments: Value...) async throws -> [Value] {
try await callAsFunction(arguments)
}
public func callAsFunction(_ arguments: [Value]) async throws -> [Value] {
try asyncify.assertNoneState()
var result = try inner(arguments)
while try asyncify.getState() == .unwinding {
guard let stopUnwind = asyncify.stopUnwindFunc else {
throw AsyncifyError.exportsNotInitialized
}
_ = try stopUnwind([])
if let storedValue = asyncify.taskStorage.getValue() {
if case .async(let task) = storedValue as! ImportResult {
let value = try await task.value
asyncify.taskStorage.setValue(value)
}
}
try asyncify.assertNoneState()
guard let startRewind = asyncify.startRewindFunc else {
throw AsyncifyError.exportsNotInitialized
}
_ = try startRewind([.i32(UInt32(Int32(bitPattern: asyncify.memoryLayout.dataAddr)))])
result = try inner(arguments)
}
try asyncify.assertNoneState()
return result
}
}
/// A collection of async-wrapped exports
public struct AsyncExports {
private let asyncify: Asyncify
private let instance: Instance
fileprivate init(asyncify: Asyncify, instance: Instance) {
self.asyncify = asyncify
self.instance = instance
}
public subscript(function name: String) -> AsyncFunction? {
guard let function = instance.exports[function: name],
!name.starts(with: "asyncify_")
else {
return nil
}
return asyncify.wrapExport(function)
}
public subscript(memory name: String) -> Memory? {
instance.exports[memory: name]
}
public subscript(table name: String) -> Table? {
instance.exports[table: name]
}
public subscript(global name: String) -> Global? {
instance.exports[global: name]
}
}
extension Module {
/// Instantiate a module with Asyncify support
public func instantiateAsync(
store: Store,
imports: Imports = [:]
) async throws -> (Instance, Asyncify) {
let asyncify = Asyncify()
let instance = try instantiate(store: store, imports: imports)
try asyncify.initialize(instance: instance)
return (instance, asyncify)
}
/// Instantiate a module with async imports
public func instantiateAsync(
store: Store,
asyncImports: AsyncImports
) async throws -> (Instance, Asyncify) {
let asyncify = Asyncify()
var wrappedImports = Imports()
for (moduleName, moduleImports) in asyncImports.definitions {
for (importName, asyncImport) in moduleImports {
let wrappedBody = asyncify.wrapImport(asyncImport.body)
let function = Function(store: store, type: asyncImport.type, body: wrappedBody)
wrappedImports.define(module: moduleName, name: importName, function)
}
}
let instance = try instantiate(store: store, imports: wrappedImports)
try asyncify.initialize(instance: instance)
return (instance, asyncify)
}
}
/// A collection of async import functions
public struct AsyncImports {
struct AsyncImport {
let type: FunctionType
let body: @Sendable (Caller, [Value]) throws -> ImportResult
}
internal var definitions: [String: [String: AsyncImport]] = [:]
public init() {}
/// Define an import function that returns ImportResult
public mutating func define(
module: String,
name: String,
type: FunctionType,
_ implementation: @escaping @Sendable (Caller, [Value]) throws -> ImportResult
) {
definitions[module, default: [:]][name] = AsyncImport(type: type, body: implementation)
}
/// Define an import function with parameter and result types
public mutating func define(
module: String,
name: String,
parameters: [ValueType],
results: [ValueType] = [],
_ implementation: @escaping @Sendable (Caller, [Value]) throws -> ImportResult
) {
define(
module: module,
name: name,
type: FunctionType(parameters: parameters, results: results),
implementation
)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment