Zenn
📗

WebAssembly と JavaScript との間で自在にデータをやりとりする

2022/11/08に公開
JavaScript
Rust
TypeScript
WebAssembly
#
serde
tech

はじめに

引き続き Rust + WebAssembly + SolidJS で遊んでいます。前回は、Rust 側で作成した文字列を JavaScript 側で console.log に出力することを考えましたが、今回は JavaScript 側から何らかのデータを Rust 側へ渡すことを考えたいと思います。今回も、wasm-bindgen[1] に頼らずにやっていきましょう。

https://zenn.dev/a24k/articles/20221012-wasmple-simple-console

メモリの確保と管理

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 と名前をつけておきます。

buffer/buffer.rs
pub type BufferPtr = usize;

#[derive(Debug)]
pub struct Buffer {
    ptr: BufferPtr,
    layout: Layout,
}

確保したバッファのサイズを取得したいケースも多々あるのですが、layout.size() で取得できそうなのでいったんこれだけあれば十分そうです。

https://doc.rust-lang.org/stable/src/core/alloc/layout.rs.html#39

メモリの確保

メモリの確保は、Layout を作成してそれを std::alloc::alloc に渡すだけなのですが、u8 や u16 など「任意の型」と「その型の値の数」を渡してメモリ確保ができると、JavaScript 側から Uint8Array や Uint16Array でアクセスする際に分かりやすい気がしたので、その様な関数にしてみます。

buffer/buffer.rs
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 から自在にアクセスすることが出来そうです。

buffer/buffer.rs
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 トレイトを実装してあげれば良さそうです。

buffer/buffer.rs
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 の情報にアクセス出来れば良い気がします。

buffer/manager.rs
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() することになります。こういう感じにしちゃいましょう。

buffer/manager.rs
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 は生きることになります。

buffer/manager.rs
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 側から利用）を用意しておきます。

buffer/manager.rs
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 されてメモリも解放されるので、こんな感じでイケそうです。いちおう全削除用の関数も用意しておきます。

buffer/manager.rs
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 トレイトを実装しておきます。

buffer/export.rs
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 を呼び出す関数を公開しておけば良さそうです。目がチカチカします。

buffer/export.rs
#[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 で書いています。型はこんな感じ。

wasm/buffer.ts
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 へのアクセスになります。良さそうです。

wasm/buffer.ts
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 との間でやりとりすることが出来る様になりました。試しに、入力されたテキストを逆順にするデモを作ってみました。

screen0001

実際に動くものはこちらから。ログを見ると、入力するたびにバッファの確保と解放がされている様子が分かるかと思います。GC じゃなくて即時解放されるの・・・いいよね。すぐ次の変換時に同じアドレスが再利用されたりします。ソースコードは、↓ からどうぞ。

https://github.com/a24k/wasmple/tree/55fc78f078a54d8e60061489c53abb8284a41605

もっとなめらかに・・・

上のデモの Rust 側のコードがコレなのですが、lock とか unwrap とか・・・どうにもゴチャゴチャしています。文字列を受け渡すくらい、もう少しなめらかに出来たいものです。

https://github.com/a24k/wasmple/blob/55fc78f078a54d8e60061489c53abb8284a41605/wasmple/src/lib.rs

Rust 上で、任意の型とメモリバッファとの変換を考えてみます。こんな感じのトレイトが実装されていると良い様な気がします。

buffer/convert.rs
pub trait BufferConverter<T> {
    fn from(ptr: BufferPtr) -> Option<T>;
    fn into(&self) -> Option<BufferPtr>;
}

これを実装しておいて、こんな ↓ 関数を用意しておけば、buffer::into::<String>(ptr) みたいな感じで変換が出来そうです。やってみましょう。

buffer/convert.rs
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 になっていることでネストを減らせるのですね。

buffer/convert/string.rs
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 側にも対になるコードを書いておきます。文字列は、頻繁にやりとりすると思われるので、大事ですね。

wasm/buffer.ts
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);
        },
    };
}

実際に動くものはこちらから。コードは、こんな ↓ 感じです。

https://github.com/a24k/wasmple/tree/583fa914389a30864332334b7d7e9491766bb451

JSON の受け渡し（untyped）

さて、JavaScript との間で文字列の受け渡しが出来るということは、その上に JSON を載せられると幸せになれる気がしますよね。Rust 上で JSON を扱う際には、serde_json[4] というクレートを使うのが良さそうです。こんな感じで BufferConverter を実装してみます。

buffer/convert/json.rs
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();

https://github.com/a24k/wasmple/tree/bf6b5a18e306a5078bf50aa1d468a730cb275ea6

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 というトレイトを用意して、自分でマークをつける様にしてみます。

buffer/convert/json.rs
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 マクロでマークを付けられる様になりました。

https://zenn.dev/link/articles/20221113-wasmple-define-macros

最終的に、パラメタも戻り値も typed-JSON になったものがコチラです。結果的に、一気に処理して返せる様になったため、WebAssembly のコール数も減っていて良いと思います。

screen0002

実際に動くものはこちらから。コードは、こんな ↓ になりました。

https://github.com/a24k/wasmple/tree/2e718be5e94ad20c467bebecb1cc74bba586f01a

おわりに

なんとなく Rust と JavaScript との間で、不自由なくデータのやりとりが出来るようになってきました。なんか常に「もっと良い組み方があるんじゃないか？」という気持ちと戦っていますが、今の私の Rust スキルではこの辺りが分相応かなとは思います。今日はここまで。

脚注

  1. wasm-bindgen - Facilitating high-level interactions between Wasm modules and JavaScript. ↩︎

  2. JavaScript Interoperation - Going Beyond Numerics ↩︎

  3. 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. ↩︎

  4. serde_json - Serde JSON provides efficient, flexible, safe ways of converting data. ↩︎

Discussion