WebAssembly 関数に JavaScript のクロージャを擬似的に渡す in Rust

Tech
Share

現状の WebAssembly の仕様では JavaScript の関数は Rust へと直接渡せないということなので、効率は悪いですが一つの回避策を考えてみました。
仲介役を JavaScript 側に立たせ、 JavaScript 側ではほぼ任意なインターフェース、 Rust 側では固定的なインターフェースにします。
現状では実用的ではないレベルですが、どうしても避けられない場合はこうしてみるのがいいかもしれません。

もっと使いやすいやり方ができるようになるといいのですが。。。

サンプルプロジェクトはこちら

作るもの

基本例としての単純なメソッド withClosureExample と、応用例としてのメソッド fold を作ります。
前者はとても簡単で、 withClosureExample((a1, a2) => a1 + a2, 2, 3) と渡すと 2 + 3 = 5 に、 withClosureExample((a1, a2) => a1 * a2, 2, 3) と渡すと 2 * 3 = 6 となるものです。
後者は内部で持つデータに対して畳み込みするものです。例えば、 [2, 3, 5] の状態で fold((l, r) => l + r, 0) と渡すと (((0 + 2) + 3) + 5) = 10 に、 fold((l, r) => l * r, 1) と渡すと (((1 * 2) * 3) * 5) = 30 となるものです。

サンプルでは、 wasm の出力のラップする JS “ラッパーライブラリ” において call_closure という仲介役を設けることで、クロージャ渡しを擬似的に再現します。

WASM 実装 - lib.rs

Rust で WebAssembly を出力する を参考にしつつ、 extern と #[no_mangle] を駆使して実装します。
call_closure については、引数を変換が必要ない i32 と f64 に固定し、その数を 4 に固定します。

lib.rs
1
2
3
4
5
6
7
8
extern "C" {
    fn call_closure(closure_id: i32, arg1: f64, arg2: f64, arg3: f64) -> f64;
}

#[no_mangle]
pub unsafe extern "C" fn with_closure_example(closure_id: i32, arg1: f64, arg2: f64, arg3: f64) -> f64 {
    call_closure(closure_id, arg1, arg2, arg3)
}

こうすると、 JS からは with_closure_examle で呼べるようになります。ただし、後述しますが call_closure の実態は指定して呼ぶ必要があります。

応用例については lib.rs を参照してください。
init_data は内部データ用の配列の初期化、 fold が上述した関数、 fold_as_sum_in_rust はベンチマーク用にすべてを Rust で実装した関数です。

ラッパー実装 - lib.js

JS から wasm 関数を利用するには、下のコードのハイライト部分の処理を行い、 results.instance.exports.??? という呼び出しを行います。
ラッパーとしては、まずパラメータを渡し、さらにこれを利用するように包んだ関数を export します。（ wasm を利用すると形はどうであれ非同期処理になります）

lib.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
const closureDict = {};

const imports = {
    env: {
        call_closure: function (closureId, ...args) {
            return closureDict[closureId](...args);
        }
    }
};

let counter = 0;

function generateId() {
    return counter++;
}

export default fetch('./target/wasm32-unknown-unknown/debug/pseudo_pass_js_closure_to_wasm_in_rust.wasm')
    .then(response => response.arrayBuffer())
    .then(bytes => WebAssembly.instantiate(bytes, imports))
    .then(results => ({
        withClosureExample: (closure, ...args) => {
            if (typeof closure === "function") {
                const id = generateId();
                closureDict[id] = closure;
                const result = results.instance.exports.with_closure_example(id, ...args);
                delete closureDict[id];
                return result;
            }
            else {
                throw new Error("`closure` is not a function!");
            }
        },
        initData: results.instance.exports.init_data,
        fold: (closure, init) => {
            if (typeof closure === "function") {
                const id = generateId();
                closureDict[id] = closure;
                const result = results.instance.exports.fold(id, init);
                delete closureDict[id];
                return result;
            }
            else {
                throw new Error("`closure` is not a function!");
            }
        },
        foldAsSumInRust: (init) => {
            return results.instance.exports.fold_as_sum_in_rust(init);
        }
    }));

一般的な wasm 利用の重要な点は、 3 行目から 5 行目と 19 行目です。
wasm 内で extern 指定したものは、こうして実装を入れてやらないと実行時にエラーが出てきます。

このサンプルでは、 call_closure という仲介役を設けると説明しました。本当に簡単な戦略ですが、呼び出し時にクロージャをストアに登録し、 wasm から call_closure を介して呼んでいます。

その他の点には面白みはあまりありません。

ラッパーの利用側 - index.html

最後に、ラッパーの実行例を示します。

index.html
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
<script type="module">
    import _lib from './lib.js';

    _lib.then(lib => {
        console.log(lib.withClosureExample(Math.sqrt, 5));
        console.log(lib.withClosureExample((a, b) => a * b, 2, 3));
        console.log(lib.withClosureExample(function (a, b) {
            console.log(arguments);
            return a + b;
        }, 2, 3));

        lib.initData();
        const data = [];
        for (let i = 0; i < 1000000; i++) {
            data[i] = i + 1;
        }

        const foldFunc = (l, r) => {
            return l + r;
        };

        console.time("wasm-with-js");
        console.log(lib.fold(foldFunc, 0));
        console.timeEnd("wasm-with-js");

        console.time("js");
        console.log(data.reduce(foldFunc, 0));
        console.timeEnd("js");

        console.time("wasm");
        console.log(lib.foldAsSumInRust(0));
        console.timeEnd("wasm");
    });
</script>

lib.withClosureExample には複数の種類のクロージャを渡せています。

lib.fold に関しては、クロージャ利用 “wasm-with-js”、全て JS 利用 “js”、全て wasm 利用 “wasm” の 3 つの動作速度を比較します。
初期化時間やデータの移動処理は時間に含めないよう書きました。

実行速度

wasm を使うからには速度が気になるということでやってみましたが、こんな感じになりました。

計算方法 実行結果 Chrome 68 64bit 速度 Firefox 61 64bit 速度
wasm-with-js 500000500000 121.224853515625ms 98ms
js 500000500000 17.569091796875ms 5ms
wasm 500000500000 75.505126953125ms 37ms

“wasm-with-js” v.s. “js” が JIT によりこのような差になるのは不思議ではないですが、 “js” v.s. “wasm” がこうなるのはちょっと想定外でした。
今回は計算対象が足し算の連続で JIT が効きやすかったからかもしれません。シミュレーションのような複雑な計算をするならば “wasm” が優勢になってくるかもしれません。

まとめ

  • wasm の提供する関数へクロージャを渡したい場合の素朴な方法を紹介しました。
  • その場合の実行速度を比較しました。