115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
from typing import Callable, List, Union
|
|
|
|
import ctypes
|
|
import struct
|
|
|
|
import wasmtime
|
|
|
|
from phasmplatform.common import valuetype
|
|
from phasmplatform.common.method import Method
|
|
from phasmplatform.common.router import BaseRouter
|
|
from phasmplatform.common.value import Value
|
|
from phasmplatform.common.valuetype import ValueType
|
|
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 value_to_wasm(self, val: Value) -> Union[None, int, float]:
|
|
if val.value_type is valuetype.bytes:
|
|
assert isinstance(val.data, bytes) # type hint
|
|
return self.alloc_bytes(val.data)
|
|
|
|
raise NotImplementedError(val)
|
|
|
|
def value_from_wasm(self, value_type: ValueType, val: Union[None, int, float]) -> Value:
|
|
if value_type is valuetype.bytes:
|
|
assert isinstance(val, int) # typ hint
|
|
return Value(valuetype.bytes, self.read_bytes(val))
|
|
|
|
raise NotImplementedError(value_type, val)
|
|
|
|
def do_call(self, method: Method, args: List[Value], on_result: Callable[[Value], None]) -> None:
|
|
wasm_method = self.exports[method.name]
|
|
assert isinstance(wasm_method, wasmtime.Func)
|
|
|
|
act_args = [self.value_to_wasm(x) for x in args]
|
|
result = wasm_method(self.store, *act_args)
|
|
assert result is None or isinstance(result, (int, float, )) # type hint
|
|
on_result(self.value_from_wasm(method.return_type, result))
|
|
|
|
# callback(UntypedValue(result)) # TODO: This returns a bytes pointer, but we can't detect that in advance
|
|
|
|
# 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)
|