WebAssembly と JavaScript との間で自在にデータをやりとりする
はじめに
引き続き Rust + WebAssembly + SolidJS で遊んでいます。前回は、Rust 側で作成した文字列を JavaScript 側で console.log に出力することを考えましたが、今回は JavaScript 側から何らかのデータを Rust 側へ渡すことを考えたいと思います。今回も、wasm-bindgen[1] に頼らずにやっていきましょう。
メモリの確保と管理
WebAssembly のメモリ空間は、シンプルな Linear Memory(線形メモリ)になっています。Rust と JavaScript との間で、「プリミティブな数値型」より大きなデータをやりとりするためには、この Linear Memory 上にデータを配置するのが良さそうです[2]。
単一のメモリバッファの確保
メモリの確保は Rust 側で実施します。線形メモリ上には WebAssembly 上で使用されるあらゆるデータが格納されるため、Rust のメモリアロケーターに管理してもらうのが吉かと思います。
Rust でメモリを確保するには以下の関数を用います。zeroed
の方はその名の通りバッファをゼロ埋めしてくれます。無印の方が多分速いですが、もう令和なので zeroed
の方を使ってみます。
pub unsafe fn alloc(layout: Layout) -> *mut u8
pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8
Layout
は、確保する領域のサイズとアラインメントを指定するための構造体です。この値は、メモリを解放する際にも必要になります。
pub unsafe fn dealloc(ptr: *mut u8, layout: Layout)
こんな様子。なるほど。もちろんどちらも unsafe
なので、うまく包んでいきたいと思います。
確保したバッファを表現する構造体
まずは、確保したバッファを表現する構造体を考えてみます。ポインタ(アドレス)と Layout
が必要になりそうです。ポインタは *mut u8
で手に入るのですが、今後の扱いでスレッドセーフじゃないと怒られる(ちょっとどうするのが最適か?は分からなかった・・・)ので usize
で扱うことにしました。この値は JavaScript とのやりとりでも多用されるため、BufferPtr
と名前をつけておきます。
pub type BufferPtr = usize;
#[derive(Debug)]
pub struct Buffer {
ptr: BufferPtr,
layout: Layout,
}
確保したバッファのサイズを取得したいケースも多々あるのですが、layout.size()
で取得できそうなのでいったんこれだけあれば十分そうです。
メモリの確保
メモリの確保は、Layout
を作成してそれを std::alloc::alloc
に渡すだけなのですが、u8
や u16
など「任意の型」と「その型の値の数」を渡してメモリ確保ができると、JavaScript 側から Uint8Array
や Uint16Array
でアクセスする際に分かりやすい気がしたので、その様な関数にしてみます。
impl Buffer {
pub fn alloc<T>(length: usize) -> Option<Self> {
if length == 0 {
return None;
}
let align = std::mem::align_of::<T>();
let unit = std::mem::size_of::<T>();
let layout = Layout::from_size_align(length * unit, align).ok()?;
let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
console::debug(format!(
"[wasm] Buffer::alloc ptr = {:?} layout = {:?}",
ptr, layout
));
match ptr.is_null() {
true => None,
false => Some(Self {
ptr: ptr as BufferPtr,
layout,
}),
}
}
}
型との紐付けに関しては、当初は Buffer<T>
みたいにして確保した際の型で縛る様な作りも考えてみたのですが、私の Rust 理解度に対して難し過ぎたのと、結局 JavaScript 側からは自由にアクセスできてしまうので不毛に感じて、このカタチになっています。
メモリへのアクセス
確保したメモリに対して Rust からアクセスする際には、スライス &[T]
として扱うのが良さそうです。生のポインタからスライスを仕立てるには、std::slice::from_raw_parts
を使用します。これで、確保済みのバッファに対して Rust から自在にアクセスすることが出来そうです。
impl Buffer {
pub fn slice<T>(&self) -> &[T] {
unsafe { std::slice::from_raw_parts(self.ptr as *const T, self.length::<T>()) }
}
pub fn slice_mut<T>(&mut self) -> &mut [T] {
unsafe { std::slice::from_raw_parts_mut(self.ptr as *mut T, self.length::<T>()) }
}
}
メモリの解放
不要になったメモリは解放しなければいけません。Buffer
が drop
される際に、解放されると良い気がします。RAII っていうヤツです。Drop
トレイトを実装してあげれば良さそうです。
impl Drop for Buffer {
fn drop(&mut self) {
unsafe { std::alloc::dealloc(self.ptr as *mut u8, self.layout) }
console::debug(format!("[wasm] dealloc ptr = {:?}", self.ptr));
}
}
複数のメモリバッファの管理
これで、単一のメモリバッファの確保・アクセス・解放が出来る様になりましたが、現実には複数のバッファを確保して Rust と JavaScript との間でやりとりする必要が出てきます。例えば 1 回の関数呼び出しにしても、複数のパラメタや戻り値をやりとりしたいケースがありそうです。ここでは、その方法について考えていきたいと思います。
JavaScript 側からのアクセスも考えると、BufferPtr
の値をキーに Buffer
の情報にアクセス出来れば良い気がします。
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
pub struct BufferManager {
ring: HashMap<BufferPtr, Arc<Mutex<Buffer>>>,
}
Buffer
を参照したあとの(Rust コード内での)取り回しを考えて、参照カウントポインタに包んでいます。BufferManager
を Singleton 風にしたいのですが、static
変数はスレッドセーフにしないと怒られる様なので Arc
と Mutex
を使っています。
BufferManager
を Singleton 風に扱いたいので、OnceCell
[3] を用いて static
変数にしています。Mutex
にも包まれているため、BufferManager
にアクセスする際には、必ず lock()
することになります。こういう感じにしちゃいましょう。
use once_cell::sync::OnceCell;
use std::sync::{Mutex, MutexGuard};
impl BufferManager {
pub fn lock() -> MutexGuard<'static, BufferManager> {
static MANAGER: OnceCell<Mutex<BufferManager>> = OnceCell::new();
MANAGER.get_or_init(|| {
Mutex::new(BufferManager { ring: HashMap::new() })
}).lock().unwrap()
}
}
Mutex
の lock()
は空くまでブロックしてくれるのだし、失敗した場合はもう panic で良いかなと思っています・・・。
メモリの確保
BufferManager
が出来たので、コイツを通じてメモリを確保する様にしましょう。Buffer
を作成して ring: HashMap
に追加すれば OK です。clone()
したものを HashMap
に入れているので、両方が解放されるまで Buffer
は生きることになります。
impl BufferManager {
pub fn alloc<T>(&mut self, length: usize) -> Option<Arc<Mutex<Buffer>>> {
let buf = Buffer::alloc::<T>(length)?;
let ptr = buf.ptr();
let arc = Arc::new(Mutex::new(buf));
self.ring.insert(ptr, arc.clone());
console::debug(format!("[wasm] dump {:?}", self));
Some(arc)
}
}
メモリへのアクセス
バッファへのアクセサ(主に Rust 側から利用)とバッファの長さを取得する関数(主に JavaScript 側から利用)を用意しておきます。
impl BufferManager {
pub fn get(&self, ptr: BufferPtr) -> Option<Arc<Mutex<Buffer>>> {
Some(self.ring.get(&ptr)?.clone())
}
pub fn length<T>(&self, ptr: BufferPtr) -> usize {
self.get(ptr)
.map_or(0, |arc| arc.lock().unwrap().length::<T>())
}
}
メモリの解放
Arc
のカウンタがゼロになれば Buffer
が drop
されてメモリも解放されるので、こんな感じでイケそうです。いちおう全削除用の関数も用意しておきます。
impl BufferManager {
pub fn dealloc(&mut self, ptr: BufferPtr) -> Option<Arc<Mutex<Buffer>>> {
let removed = self.ring.remove(&ptr);
console::debug(format!("[wasm] dump {:?}", self));
removed
}
pub fn clear(&mut self) {
self.ring.clear();
console::debug(format!("[wasm] dump {:?}", self));
}
}
JavaScript からのアクセス
Rust 内でのバッファの確保は大体できたので、JavaScript 側からのアクセスについて考えていきます。JavaScript 側からメモリバッファにアクセスする際には、WebAssembly からエクスポートされた ArrayBuffer
に対して TypedArray
を通じてアクセスすることが多いと思うので、それを目指して仕立てていきます。
Rust 側からのエクスポート
任意のバッファに対して型を指定してアクセスしたいのですが、JavaScript と Rust との間で型情報をやりとりする方法が多分ないので enum
で整数値にマッピングして扱う様にします。enum
から数値は as
で変換出来るのですが、逆は不可なので From
トレイトを実装しておきます。
enum T {
I8, U8, I16, U16, I32, U32, I64, U64, F32, F64,
}
impl From<u8> for T {
fn from(t: u8) -> Self {
match t {
0 => T::I8, 1 => T::U8, 2 => T::I16, 3 => T::U16,
4 => T::I32, 5 => T::U32, 6 => T::I64, 7 => T::U64,
8 => T::F32, 9 => T::F64, _ => T::U8,
}
}
}
あとはこんな感じで、BufferManager
を呼び出す関数を公開しておけば良さそうです。目がチカチカします。
#[no_mangle]
pub extern "C" fn buffer_alloc(t: u8, len: usize) -> BufferPtr {
let arc = match T::from(t) {
T::I8 => BufferManager::lock().alloc::<i8>(len),
T::U8 => BufferManager::lock().alloc::<u8>(len),
T::I16 => BufferManager::lock().alloc::<i16>(len),
T::U16 => BufferManager::lock().alloc::<u16>(len),
T::I32 => BufferManager::lock().alloc::<i32>(len),
T::U32 => BufferManager::lock().alloc::<u32>(len),
T::I64 => BufferManager::lock().alloc::<i64>(len),
T::U64 => BufferManager::lock().alloc::<u64>(len),
T::F32 => BufferManager::lock().alloc::<f32>(len),
T::F64 => BufferManager::lock().alloc::<f64>(len),
};
arc.map_or(0, |arc| arc.lock().unwrap().ptr())
}
#[no_mangle]
pub extern "C" fn buffer_length(t: u8, ptr: BufferPtr) -> usize {
match T::from(t) {
T::I8 => BufferManager::lock().length::<i8>(ptr),
T::U8 => BufferManager::lock().length::<u8>(ptr),
T::I16 => BufferManager::lock().length::<i16>(ptr),
T::U16 => BufferManager::lock().length::<u16>(ptr),
T::I32 => BufferManager::lock().length::<i32>(ptr),
T::U32 => BufferManager::lock().length::<u32>(ptr),
T::I64 => BufferManager::lock().length::<i64>(ptr),
T::U64 => BufferManager::lock().length::<u64>(ptr),
T::F32 => BufferManager::lock().length::<f32>(ptr),
T::F64 => BufferManager::lock().length::<f64>(ptr),
}
}
#[no_mangle]
pub extern "C" fn buffer_dealloc(ptr: BufferPtr) {
BufferManager::lock().dealloc(ptr);
}
#[no_mangle]
pub extern "C" fn buffer_clear() {
BufferManager::lock().clear();
}
JavaScript 側からの利用
enum
を使いたかったこともあり、すべて TypeScript で書いています。型はこんな感じ。
export enum Type {
I8, U8, I16, U16, I32, U32, I64, U64, F32, F64,
}
export type BufferPtr = number;
type FnAlloc = (t: Type, len: number) => number;
type FnLength = (t: Type, ptr: BufferPtr) => number;
type FnDealloc = (ptr: BufferPtr) => void;
type FnClear = () => void;
WasmBuffer
というクラスを作成して、BufferManager
とのやりとりをまとめちゃいます。Rust でいうスライスへのアクセスに相当するのが、Int8Array
などの TypedArray
へのアクセスになります。良さそうです。
export class WasmBuffer {
private memory: WebAssembly.Memory;
public alloc: FnAlloc;
public length: FnLength;
public dealloc: FnDealloc;
public clear: FnClear;
constructor(wasm: WebAssembly.Exports) {
this.memory = wasm.memory as WebAssembly.Memory;
this.alloc = wasm.buffer_alloc as FnAlloc;
this.length = wasm.buffer_length as FnLength;
this.dealloc = wasm.buffer_dealloc as FnDealloc;
this.clear = wasm.buffer_clear as FnClear;
}
public slice = {
i8: (ptr: BufferPtr): Int8Array => { return new Int8Array(this.memory.buffer, ptr, this.length(Type.I8, ptr)); },
u8: (ptr: BufferPtr): Uint8Array => { return new Uint8Array(this.memory.buffer, ptr, this.length(Type.U8, ptr)); },
i16: (ptr: BufferPtr): Int16Array => { return new Int16Array(this.memory.buffer, ptr, this.length(Type.I16, ptr)); },
u16: (ptr: BufferPtr): Uint16Array => { return new Uint16Array(this.memory.buffer, ptr, this.length(Type.U16, ptr)); },
i32: (ptr: BufferPtr): Int32Array => { return new Int32Array(this.memory.buffer, ptr, this.length(Type.I32, ptr)); },
u32: (ptr: BufferPtr): Uint32Array => { return new Uint32Array(this.memory.buffer, ptr, this.length(Type.U32, ptr)); },
i64: (ptr: BufferPtr): BigInt64Array => { return new BigInt64Array(this.memory.buffer, ptr, this.length(Type.I64, ptr)); },
u64: (ptr: BufferPtr): BigUint64Array => { return new BigUint64Array(this.memory.buffer, ptr, this.length(Type.U64, ptr)); },
f32: (ptr: BufferPtr): Float32Array => { return new Float32Array(this.memory.buffer, ptr, this.length(Type.F32, ptr)); },
f64: (ptr: BufferPtr): Float64Array => { return new Float64Array(this.memory.buffer, ptr, this.length(Type.F64, ptr)); },
};
}
これで、任意のバッファを Rust と JavaScript との間でやりとりすることが出来る様になりました。試しに、入力されたテキストを逆順にするデモを作ってみました。
実際に動くものはこちらから。ログを見ると、入力するたびにバッファの確保と解放がされている様子が分かるかと思います。GC じゃなくて即時解放されるの・・・いいよね。すぐ次の変換時に同じアドレスが再利用されたりします。ソースコードは、↓ からどうぞ。
もっとなめらかに・・・
上のデモの Rust 側のコードがコレなのですが、lock
とか unwrap
とか・・・どうにもゴチャゴチャしています。文字列を受け渡すくらい、もう少しなめらかに出来たいものです。
Rust 上で、任意の型とメモリバッファとの変換を考えてみます。こんな感じのトレイトが実装されていると良い様な気がします。
pub trait BufferConverter<T> {
fn from(ptr: BufferPtr) -> Option<T>;
fn into(&self) -> Option<BufferPtr>;
}
これを実装しておいて、こんな ↓ 関数を用意しておけば、buffer::into::<String>(ptr)
みたいな感じで変換が出来そうです。やってみましょう。
pub fn into<T>(ptr: BufferPtr) -> Option<T>
where
T: BufferConverter<T>,
{
T::from(ptr)
}
pub fn from<T>(obj: T) -> Option<BufferPtr>
where
T: BufferConverter<T>,
{
T::into(&obj)
}
文字列の受け渡し
String
型との変換はこんな感じになります。戻り値が Option
になっていることでネストを減らせるのですね。
impl BufferConverter<String> for String {
fn from(ptr: BufferPtr) -> Option<String> {
let arc = BufferManager::lock().get(ptr)?;
let buf = arc.lock().unwrap();
let slice = buf.slice::<u16>();
String::from_utf16(slice).ok()
}
fn into(&self) -> Option<BufferPtr> {
let utf16: Vec<u16> = self.encode_utf16().collect();
let arc = BufferManager::lock().alloc::<u16>(utf16.len())?;
let mut buf = arc.lock().unwrap();
let slice = buf.slice_mut::<u16>();
slice.copy_from_slice(&utf16[..]);
Some(buf.ptr())
}
}
JavaScript 側にも対になるコードを書いておきます。文字列は、頻繁にやりとりすると思われるので、大事ですね。
export class WasmBuffer {
public from = {
string: (str: string): BufferPtr => {
const len = str.length; // number of UTF-16 code units
const ptr = this.alloc(Type.U16, len);
const buf = this.slice.u16(ptr);
for (let i = 0; i < len; ++i) { buf[i] = str.charCodeAt(i); }
return ptr;
},
};
public to = {
string: (ptr: BufferPtr): string => {
const chars = this.slice.u16(ptr);
return String.fromCharCode(...chars);
},
};
}
実際に動くものはこちらから。コードは、こんな ↓ 感じです。
JSON の受け渡し(untyped)
さて、JavaScript との間で文字列の受け渡しが出来るということは、その上に JSON を載せられると幸せになれる気がしますよね。Rust 上で JSON を扱う際には、serde_json[4] というクレートを使うのが良さそうです。こんな感じで BufferConverter
を実装してみます。
use serde_json::Value;
impl BufferConverter<Value> for Value {
fn from(ptr: BufferPtr) -> Option<Value> {
serde_json::from_str(&super::into::<String>(ptr)?).ok()
}
fn into(&self) -> Option<BufferPtr> {
super::from(serde_json::to_string(self).ok()?)
}
}
これで、こんな具合にバッファ上のデータを JSON として扱える様になります。べんり。
let input_json: Value = buffer::into(input_ptr)?;
let input_str_a = input_json["a"].clone();
let input_str_b = input_json["b"].clone();
JSON の受け渡し(typed)
まぁまぁ良いのですが、Rust と TypeScript とのやりとりなのに型があやふやですね。なんとかしましょう。serde_json 上で typed な変換を行うには、こんな感じの構造体を定義します。
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize)]
struct FnConvertParameters {
a: String,
b: String,
}
serde::Serialize
と serde::Deserialize
というトレイトが付くので、これらに対して BufferConverter
を実装すれば良いかと思ったのですが、これらは String
とかに対しても実装されているようで、ちょっと広過ぎます。仕方がないので、JsonConvertee
というトレイトを用意して、自分でマークをつける様にしてみます。
pub trait JsonConvertee: serde::Serialize + serde::de::DeserializeOwned {}
impl JsonConvertee for serde_json::Value {}
impl<T> BufferConverter<T> for T
where
T: JsonConvertee,
{
fn from(ptr: BufferPtr) -> Option<T> {
serde_json::from_str(&super::into::<String>(ptr)?).ok()
}
fn into(&self) -> Option<BufferPtr> {
super::from(serde_json::to_string(self).ok()?)
}
}
JsonConvertee
を実装することで、任意の構造体に対して typed な JSON 変換が出来る様になりました。また、serde_json::Value
に対しても JsonConvertee
マークを付けてあるため、untyped な場合も含めて同じコードで処理される様になっています。
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize)]
struct FnConvertParameters {
a: String,
b: String,
}
impl JsonConvertee for FnConvertParameters {}
2022-11-15 追記)derive マクロでマークを付けられる様になりました。
最終的に、パラメタも戻り値も typed-JSON になったものがコチラです。結果的に、一気に処理して返せる様になったため、WebAssembly のコール数も減っていて良いと思います。
実際に動くものはこちらから。コードは、こんな ↓ になりました。
おわりに
なんとなく Rust と JavaScript との間で、不自由なくデータのやりとりが出来るようになってきました。なんか常に「もっと良い組み方があるんじゃないか?」という気持ちと戦っていますが、今の私の Rust スキルではこの辺りが分相応かなとは思います。今日はここまで。
-
wasm-bindgen - Facilitating high-level interactions between Wasm modules and JavaScript. ↩︎
-
once_cell - once_cell provides two new cell-like types, unsync::OnceCell and sync::OnceCell. A OnceCell might store arbitrary non-Copy types, can be assigned to at most once and provides direct access to the stored contents. ↩︎
-
serde_json - Serde JSON provides efficient, flexible, safe ways of converting data. ↩︎
Discussion