93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
from typing import Any, Callable, List, Union
|
|
|
|
import ctypes
|
|
import struct
|
|
|
|
import wasmtime
|
|
|
|
from phasmplatform.common.router import BaseRouter
|
|
from phasmplatform.common.value import BaseValue, BytesValue, UntypedValue
|
|
from .base import BaseRunner
|
|
|
|
|
|
class WasmTimeRunner(BaseRunner):
|
|
__slots__ = ('store', 'module', 'instance', 'exports')
|
|
|
|
def __init__(self, router: BaseRouter, wasm_bin: bytes) -> None:
|
|
super().__init__(router)
|
|
|
|
self.store = wasmtime.Store()
|
|
|
|
post_message_bind = wasmtime.Func(self.store, wasmtime.FuncType([
|
|
wasmtime.ValType.i32(),
|
|
wasmtime.ValType.i32(),
|
|
wasmtime.ValType.i32(),
|
|
wasmtime.ValType.i32(),
|
|
], []), self.post_message)
|
|
|
|
self.module = wasmtime.Module(self.store.engine, wasm_bin)
|
|
self.instance = wasmtime.Instance(self.store, self.module, [
|
|
post_message_bind,
|
|
])
|
|
|
|
self.exports = self.instance.exports(self.store)
|
|
|
|
def alloc_bytes(self, data: bytes) -> int:
|
|
memory = self.exports['memory']
|
|
assert isinstance(memory, wasmtime.Memory) # type hint
|
|
|
|
data_ptr = memory.data_ptr(self.store)
|
|
data_len = memory.data_len(self.store)
|
|
|
|
alloc_bytes = self.exports['stdlib.types.__alloc_bytes__']
|
|
assert isinstance(alloc_bytes, wasmtime.Func)
|
|
|
|
ptr = alloc_bytes(self.store, len(data))
|
|
assert isinstance(ptr, int) # type hint
|
|
|
|
idx = ptr + 4 # Skip the header from header from __alloc_bytes__
|
|
for byt in data:
|
|
assert idx < data_len
|
|
data_ptr[idx] = ctypes.c_ubyte(byt)
|
|
idx += 1
|
|
|
|
return ptr
|
|
|
|
def read_bytes(self, ptr: int) -> bytes:
|
|
memory = self.exports['memory']
|
|
assert isinstance(memory, wasmtime.Memory) # type hint
|
|
|
|
data_ptr = memory.data_ptr(self.store)
|
|
data_len = memory.data_len(self.store)
|
|
|
|
raw = ctypes.string_at(data_ptr, data_len)
|
|
|
|
length, = struct.unpack('<I', raw[ptr:ptr + 4]) # Header prefixed by __alloc_bytes__
|
|
|
|
return raw[ptr + 4:ptr + 4 + length]
|
|
|
|
def convert_value(self, val: BaseValue[Any]) -> Union[int, float]:
|
|
if isinstance(val, BytesValue):
|
|
return self.alloc_bytes(val.data)
|
|
|
|
raise NotImplementedError(val)
|
|
|
|
def do_call(self, method_name: str, args: List[BaseValue[Any]], callback: Callable[[BaseValue[Any]], None]) -> None:
|
|
method = self.exports[method_name]
|
|
assert isinstance(method, wasmtime.Func)
|
|
|
|
act_args = [self.convert_value(x) for x in args]
|
|
result = method(self.store, *act_args)
|
|
|
|
callback(UntypedValue(result)) # TODO: This returns a bytes pointer, but we can't detect that in advance
|
|
|
|
def handle_message(self, namespace: bytes, topic: bytes, kind: bytes, body: bytes) -> None:
|
|
namespace_ptr = self.alloc_bytes(namespace)
|
|
topic_ptr = self.alloc_bytes(topic)
|
|
kind_ptr = self.alloc_bytes(kind)
|
|
body_ptr = self.alloc_bytes(body)
|
|
|
|
handle_message = self.exports['handle_message']
|
|
assert isinstance(handle_message, wasmtime.Func)
|
|
handle_message(self.store, namespace_ptr, topic_ptr, kind_ptr, body_ptr)
|