Rustで簡単なWebAssemblyランタイムを作ってみる

こんにちは。株式会社LabBaseでバックエンド中心のRust開発を担当している、ヤーノシュと申します。

弊社の今年のアドベントカレンダーの一環として、今一番気になっている技術であるWebAssemblyについて投稿しています。簡単なRust実装を利用して、WebAssemblyの言語としての基本的な概念を紹介したいと思います。

WAを一から勉強する傍らで書いた記事になりますので間違っていると思われる内容を発見された場合はお声をおかけいただけると幸いです。よろしくお願いします！

## WebAssemblyとは

WebAssemblyは、ウェブアプリケーションでの計算が重い処理を、JavaScriptよりパーフォマンスの良い言語で作れるようにするための、バイトコード言語です。元々は、ウェブアプリケーションの一部の処理をWebAssemblyに任せ（データ凝縮など）、アプリケーションの根本的な処理はJavaScriptで実現する想定でウェブ規格として設けられたのですが、近年ではウェブアプリケーションをRustなどといった言語で完結させようとするフレームワークも公開されつつあります。弊社でもleptos、dioxusなどといったフレームワークを、実験ベースで取り入れてみたりしています。

ただ、元々の、ブラウザーの中での実効に加え、以下のような属性からサーバーバイナリーの媒体、プラッグインシステムのプラッグインコードの媒体としても使用され始めています：

- パーフォマンスを重視して作られているため、ネイティブと近い効率の良さが見込まれる
- ブラウザーの中での実装を想定されているので安全性・組み込まれている隔離性能が高く
- 細かくアクセス制限もできたりする（Pythonなどといったスクリプト言語を組み込むよりはずっと安全に）
    - 許可なしではランダム値の生成するできない
    - ネットワーク、ファイルシステムへのアクセスは然り
- WASIで特定のOSへの依存が解消され、Javaと同様に、一つのバイナリーをどんなハードウェアでも実行できる世界が見えてきている
- 公開スタンダードなため、各自好きな言語（自分はRustになりますが）で作れていて、それらの間の通信・データ交換も容易くなる

他の使い道として、Rustで使われる「proc macro」をより安全に実効できる（あるいは、コンパイルされた状態で共有できる）、などといった提案も出回っています。

## 今回の実装のスコープ

あくまでも理解を深めるための実装になりますので一部のインストラクションしか実装せず、以下のようなことは一切考えないで書いているのでご了承ください：

- WAのバイナリー、あるいはテキストフォーマットのパース
- 綺麗なエラー対応
- 元々安全性を保つために必須なバリデーション
- テーブル（存在意義・使い道については少し解説しますが）
- i32以外の型（i64, f32, f64, ref）
- ブロック（if・loopは元々やりたかったが、時間切れとなり次回へと後回しします・・・）

## 概念の紹介

コードを見せながらそれぞれの概念について説明していきます。

### モージュル（の宣言と実現）

他の言語でもよくでてくる、コードをまとめる単位。WAでは関数、グローバル変数、メモリー、エクスポート、インポートなどがありますが、今回は以下のような形で実装しています：

```rust
#[derive(Debug)]
struct ModuleDeclaration {
    functions: Vec<Arc<FunctionDeclaration>>,
}

#[derive(Debug)]
struct ModuleInstance {
    memory: Vec<u8>,
    declaration: Arc<ModuleDeclaration>,
}
```

はい、関数しかありません。外の言語（ブラウザーではJavaScriptですね）と繋げることはないのでランタイムレベルで関数をインデクスで簡単に呼び出せるようにします。また、最近追加された複数メモリー（説明は後で）の使用を可能とする変更はあえて入れていないので基本的にメモリーが一個用意される実装になります。

### 関数（の宣言）

WebAssemblyには他の言語と同じように関数があり、私たちの実装では以下のような形を取ります：

```rust
#[derive(Debug)]
struct FunctionDeclaration {
    parameters: Vec<TypeDeclaration>,
    locals: Vec<TypeDeclaration>,
    return_value: Option<TypeDeclaration>,
    instructions: Vec<Instruction>,
    label: Option<String>,
}
```

まず、引数とローカル変数を事前に宣言する必要があります。インストラクションでは両方ともローカル変数としてアクセスが可能で、読み込みも書き込みもできます。

さらに任意の戻り値があります。最近のバージョンでは複数値を返せるようになりましたが、時短のため一個に制限します。

最後にはインストラクションの配列と任意のラベルがあります。

### 型について

以前説明した通り、今回は`i32`以外の型は除外します：

```rust
#[derive(Debug, Clone, Copy)]
pub enum TypeDeclaration {
    // I64,
    I32,
    // F64,
    // F32,
}
impl TypeDeclaration {
    fn default_value(&self) -> ValueType {
        match self {
            TypeDeclaration::I32 => ValueType::I32(0),
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ValueType {
    // I64(i64),
    I32(i32),
    // F64(f64),
    // F32(f32),
}
```

### スタックについて

WebAssemblyの面白い要素の一つは、ハードウェア向けのAssembly言語と違って、レジスターを利用せず、以前少し触れたローカル変数と「スタック」でデータを管理・保管しています。スタックはその通りで様々なデータが入っています（上記の値に加え関数が呼び出されたことを記録する「フレーム」や、ブロックを実装するなら「ラベル」も追加されます。なお、こちらはモデルであり、パーフォマンスを重視した実装では別々に保管することになるかもしれません、簡略化のため以下のように混ぜた実装になりました。

```rust
#[derive(Debug)]
enum StackEntry {
    Value(ValueType),
    Function(FunctionFrame),
}

#[derive(Debug)]
struct FunctionFrame {
    declaration: Arc<FunctionDeclaration>,
}

type Stack = Vec<StackEntry>;
```

ほとんどのインストラクションは、このスタックからデータを取ったり（ポップ）、追加したり（プッシュ）するのが役目です。関数を呼び出す時にも、その関数の引数をスタックで渡し、戻り値をスタックで返されます。

### メモリー

「最上」のデータしかアクセスできないことのスタックに対して、自由にインデクスでアクセスできる、バイト単位でデータを保管するメモリーがあります。今後、モージュルで複数のメモリーを同時に利用することは可能になる予定ですが、こちらも時短のため無視しています。実際のRustなどからのトランスパイルの際にヒープデータはもちろん、一部のスタックデータもメモリーのなかで保存されます。

### テーブルについて

メモリー以外に、モージュル単位でテーブルというものも作成できます。テーブルの役割は、refの管理であり、実効時にダイナミックに呼び出す関数を変えたりするなどといった、複雑なユースケースを可能にしてくれています（Rustの`&dyn Trait`などで使用）。今回はスコープ外とします。

## 実際のインストラクションの実相の解説

今回は以下のインストラクションを実装してみました：

```rust
#[derive(Debug)]
enum Instruction {
    Nop,
    Call(usize),

    // スタック操作、計算
    Const(ValueType),
    Add(TypeDeclaration),
    Sub(TypeDeclaration),
    Eq(TypeDeclaration),
    Drop,

    // 引数・ローカル変数の管理
    GetLocal(usize),
    SetLocal(usize),

    // メモリー関連
    Load(TypeDeclaration),
    Store(TypeDeclaration),
    MemorySize,
    MemoryGrow,
}
```

これから一個ずつ解説していきます・・・

### `nop`

名前の通りで、何もしないインストラクションです。

```rust
impl Instruction {
    pub fn execute(
        &self,
        stack: &mut Stack,
        module: &mut ModuleInstance,
        locals: &mut [ValueType],
    ) {
        match self {
            ..
            Instruction::Nop => {}
            .. 
        }
}
```

### `call` 

こちらは結構肝心なインストラクションになります。何なら`call`がないとプログラムの実行を会誌できないからです。まず実装を：

```rust
impl Instruction {
    pub fn execute(
        &self,
        stack: &mut Stack,
        module: &mut ModuleInstance,
        locals: &mut [ValueType],
    ) {
        match self {
            ..
            Instruction::Call(function_id) => {
                println!("Calling function id={function_id}..");

                let declaration = Arc::clone(&module.declaration.functions[*function_id]);

                let mut locals = declaration
                    .parameters
                    .iter()
                    .map(|t| {
                        // TODO: type checking
                        let Some(StackEntry::Value(v)) = stack.pop() else {
                            panic!(
                                "Tried to invoke function, \
                                but could not find parameter of type {t:?}!"
                            );
                        };

                        v
                    })
                    .rev()
                    .collect::<Vec<_>>();

                for local in &declaration.locals {
                    locals.push(local.default_value());
                }

                stack.push(StackEntry::Function(FunctionFrame {
                    declaration: Arc::clone(&declaration),
                }));

                for instruction in &declaration.instructions {
                    println!("Running instruction: {instruction:?}");
                    instruction.execute(stack, module, &mut locals[..]);
                    println!("Completed instruction, stack state: {stack:?}");
                }

                let ret_val = declaration.return_value.map(|value| {
                    // TODO: check for types!
                    let Some(v @ StackEntry::Value(_)) = stack.pop() else {
                        panic!("Function is expected to return a value of type {value:?}");
                    };

                    v
                });

                let Some(StackEntry::Function(FunctionFrame { .. })) = stack.pop() else {
                    panic!("function frame must stil be there..");
                };

                if let Some(ret_val) = ret_val {
                    stack.push(ret_val);
                }

                println!("Finished function id={function_id}..");
            }
        }
        ..
    }
}
```

まず、呼び出された関数の宣言（定義）を取得します。バリデーションをスキップしていますが、バリデーションがある場合は、すでに該当関数（インデクスで指定します）が存在していることが確認されています。

次に引数をスタックから取得し、さらにローカル変数を初期化します。アクセスする時は両方とも「ローカル変数になるので」一緒の`Vec<_>`に保存します。引数の順番を入れ替えているのも重要です！

続いては関数が呼び出されたことを記録するために「フレーム」をスタックにプッシュします。我々の実装では、こちらの主な目的は、関数に関数を呼び出した関数のスタックデータを弄らせないためです。関数が終了してからのスタック状態を確認する手段としても利用します。

そして実際に関数のインストラクションを、一個ずつ実行します。デバッグのためにスタックの状態などもログします。

インストラクションたちが無事終了したら関数の宣言に戻り値があった場合だけ、それをスタックから取得します。「フレーム」をポップし、戻り値をスタックに戻して終了します。

### スタック操作・計算系

こちらの関数ではスタックに保存されている値を弄ります：

```rust
impl Instruction {
    pub fn execute(
        &self,
        stack: &mut Stack,
        module: &mut ModuleInstance,
        locals: &mut [ValueType],
    ) {
        match self {
            .. 
            Instruction::Const(value) => stack.push(StackEntry::Value(*value)),
            Instruction::Add(t) => {
                let (
                    Some(StackEntry::Value(ValueType::I32(c2))),
                    Some(StackEntry::Value(ValueType::I32(c1))),
                ) = (stack.pop(), stack.pop())
                else {
                    panic!("Stack must contain two operands of type {t:?}");
                };

                stack.push(StackEntry::Value(ValueType::I32(c1 + c2)))
            }
            Instruction::Sub(t) => {
                let (
                    Some(StackEntry::Value(ValueType::I32(c2))),
                    Some(StackEntry::Value(ValueType::I32(c1))),
                ) = (stack.pop(), stack.pop())
                else {
                    panic!("Stack must contain two operands of type {t:?}");
                };

                stack.push(StackEntry::Value(ValueType::I32(c1 - c2)))
            }
            Instruction::Eq(t) => {
                let (
                    Some(StackEntry::Value(ValueType::I32(c2))),
                    Some(StackEntry::Value(ValueType::I32(c1))),
                ) = (stack.pop(), stack.pop())
                else {
                    panic!("Stack must contain two operands of type {t:?}");
                };

                stack.push(StackEntry::Value(ValueType::I32((c2 == c1) as i32)))
            }
            Instruction::Drop => {
                let Some(StackEntry::Value(_)) = stack.pop() else {
                    panic!("Illegal drop, can only drop data frames.");
                };
            }
            ..
        }
    }
}

```

`i32.const n`などは、固定値の`n`をスタックにプッシュします。その対比として、`drop`では最上の値を削除します。関数から戻るまえに戻り値以外のスタックのデータを全部消す必要があるので結構重要なインストラクションになります。

`add`、`sub`、`eq`ではそれぞれの数学的オペレーションを実施します。特に`sub`は順番が大事なので規格を細かく読む必要があります。

### 引数・ローカル変数操作系

こちらの二つのインストラクションでは、ローカル変数からスタックへ、スタックからローカル変数へデータを移動させます。

```rust
impl Instruction {
    pub fn execute(
        &self,
        stack: &mut Stack,
        module: &mut ModuleInstance,
        locals: &mut [ValueType],
    ) {
        match self {
            Instruction::GetLocal(idx) => {
                stack.push(StackEntry::Value(locals[*idx]));
            }
            Instruction::SetLocal(idx) => {
                let Some(StackEntry::Value(value)) = stack.pop() else {
                    panic!("Cannot push non-existing value to local variables.")
                };

                // TODO: From looking at the spec, it does not seem like there is type-checking
                //       here.. Which seems strange at least.

                locals[*idx] = value;
            }
        }
    }
}

```

コメントにも書いているのですが、規格書ではこちらの型の確認にたいしての記述がなかったのが少し不思議に感じました（読み間違いの可能性もありますが・・・　バリデーションの部分は軽くしか読んでいません）。

### メモリー操作系

スタックからメモリーへ、メモリーからスタックへのデータ移動を担うインストラクションです。

```rust
const PAGE_SIZE: usize = 65536;

impl Instruction {
    pub fn execute(
        &self,
        stack: &mut Stack,
        module: &mut ModuleInstance,
        locals: &mut [ValueType],
    ) {
        match self {
            Instruction::Load(_) => {
                // TODO: other types
                let Some(StackEntry::Value(ValueType::I32(addr))) = stack.pop() else {
                    panic!("load instruction requires address as top of stack value");
                };

                assert!(addr >= 0, "Cannot access negative addresses.");

                let addr = addr as usize;

                if module.memory.len() < addr + size_of::<i32>() {
                    panic!("access to invalid memory address attempted");
                }

                // TODO: endianness??
                stack.push(StackEntry::Value(ValueType::I32(i32::from_be_bytes(
                    module.memory[addr..addr + size_of::<i32>()]
                        .try_into()
                        .expect("bounds are correct"),
                ))));
            }
            Instruction::Store(_) => {
                // TODO: different types
                let Some(StackEntry::Value(ValueType::I32(value))) = stack.pop() else {
                    panic!("store instruction requires value to be stored on top of the stack")
                };

                // this one is always i32
                let Some(StackEntry::Value(ValueType::I32(addr))) = stack.pop() else {
                    panic!("store instruction requires value to be stored on top of the stack")
                };

                assert!(addr >= 0, "negative memory addresses are not allowed");

                let addr = addr as usize;

                if module.memory.len() < addr + size_of_val(&value) {
                    panic!("memory does not have enough capacity to store value");
                }

                module.memory[addr..addr + size_of_val(&value)]
                    .swap_with_slice(&mut value.to_be_bytes());
            }
            Instruction::MemorySize => stack.push(StackEntry::Value(ValueType::I32(
                (module.memory.len() / PAGE_SIZE) as i32,
            ))),
            Instruction::MemoryGrow => {
                let Some(StackEntry::Value(ValueType::I32(page_count))) = stack.pop() else {
                    panic!("must specify the amount of pages as i32 on top of the stack");
                };

                assert!(page_count >= 0, "cannot reduce page count");

                let new_len = module.memory.len() + (page_count as usize * PAGE_SIZE);

                if new_len > i32::MAX as usize {
                    panic!("cannot represent memory that big");
                }

                module.memory.resize(new_len, 0);
            }
        }
    }
}
```

`memory.size`でメモリーの現在の容量を取得し、`memory.grow`で追加の容量を要求します。両方とも固定のページサイズ単位で動いています。なお、メモリーの現在容量を越えたしまった場合は、「トラップ」に嵌り、実行が中断されます。我々の実装ではとりあえず`panic()`を起こして実行を中止します。

`memory.load`と`memory.store`ではデータを保存したり、取り出したりします。それぞれ`i32`のアドレスを使っており、WAの中でもすでに`i64`にアップグレードすべきではないかという議論が出ています。今回の実装では、`i32`では表せないメモリーサイズにまで伸びた時に、とりあえず`panic()`します。

## 積み残し

今回は、元々アロケーター（メモリーの部分的な割り当てを実施するシステム、あるいは関数）を実装したかったのですが、`if .. else .. end`などを時間制限で結局実装できなかったので実装も書けていない状態です。

さらに、いろんな近道を使っているのでマクロなどを使った他の型（i64, f64, f32）も対応し、外部クレートを利用して実際のWASMファイルを読み込めるようにしていきたいです。

安全性には重要なバリデーションを完全にスキップしているのもよくないと思っていますのでそこもキャッチアップしていきます。

今回の記事を少しでも楽しんでいただけたのなら光栄です。最後まで読んでくださりありがとうございます！

質問、訂正などは、遠慮なくお声をおかけください。

## 参考資料

- Mozillaのドキュメンテーション：https://developer.mozilla.org/en-US/docs/WebAssembly/Concepts
- 実際の規格書：https://webassembly.github.io/spec/core/appendix/index-instructions.html
- 今回の実装：https://github.com/jreppnow/fusm
- Rustウェブフレームワークたち：
    - https://dioxuslabs.com/　
    - https://leptos.dev/
- RustマクロをWAで実行しようプロジェクト：https://github.com/dtolnay/watt
- プラッグインをWAで実装させるエディター：