RustでJITコンパイルする電卓を実装してみる

# はじめに

JITコンパイルってなんでしょうか？インタプリタを高速化するためにネイティブコードにコンパイルしてから実行する、みたいな説明をよく見ますが、分かるようでよく分かりません。書いてみたら分かるかもしれないと思いやってみました。

念のため注意ですが、電卓程度をJITコンパイルで実装してもたぶんいいことはありません。検証とかはしてませんが、なんなら普通に作るより遅くなる可能性もあると思います。そんな感覚で読んで頂ければと思います。


# インタプリタによる実装

JITコンパイラで実装する前に、素直にインタプリタで電卓を実装してみます。完成形はこんな感じになります。

```
> 1 + 1
2
> 1 + 2 * 3
7
```

結構よくあるプログラムだと思うので、さらっと解説していきます。


### トークンへの変換

入力される文字列は数字、記号、空白などが入り混じっており、数字もつながりが表現されていない状態です。`1 + 23`だと`['1', ' ', '+', ' ', '2', '3']`みたいな感じになり、これでは後に続く処理で不便なので扱いやすい形に成形しておきます。

```rust
enum TokenKind {
    Number(u8),
    Plus,
    Minus,
    Asterisk,
    Slash,
}
```

このような`TokenKind`を定義し`Vec<TokenKind>`を生成します。この際、スペースなどの余分な文字は無視したり、予期しない文字が入力された場合はエラーにしたりします。すると上記は`[Number(1), Plus, Number(23)]`というすっきりとした形に変換されます。


### ASTへの変換

トークンに変換したら次は、抽象構文木（Abstract Syntax Tree: AST）への変換を行います。

演算子には優先度があるので、例えば`1 + 2 * 3`みたいな式を先頭から評価してしまうと先に足し算が行われてしまい、正しい結果が得られません。そこで優先度を考慮した木構造として式を表現します。

```rust
enum BinOpKind {
    Add,
    Sub,
    Mul,
    Div,
}

enum NodeKind {
    Number(u8),
    BinOp {
        kind: BinOpKind,
        lhs: Box<NodeKind>,
        rhs: Box<NodeKind>,
    },
}
```

木構造の`Node`は数値と2項演算子の2種類とし、2項演算子には加算、減算、乗算、除算の4種類があるとします。さらに2項演算子は左辺と右辺の情報をもちます。この定義を使って`1 + 2 * 3`を表現すると、ちょっと長いですがこんな感じになります。

```
BinOp {
    kind: Add,
    lhs: Number(1),
    rhs: BinOp {
        kind: Mul,
        lhs: Number(2),
        rhs: Number(3)
    }
}
```

この表現を内側から評価していけば、正しい順序で計算することができます。


### 式の評価

ASTを評価する関数を再帰的に呼ぶことで、内側から評価することができます。

```rust
fn eval(ast: NodeKind) -> u8 {
    match ast {
        Number(n) => n,
        BinOp { kind, lhs, rhs } => {
            match kind {
                Add => eval(*lhs) + eval(*rhs),
                Sub => eval(*lhs) - eval(*rhs),
                Mul => eval(*lhs) * eval(*rhs),
                Div => eval(*lhs) / eval(*rhs),
            }
        },
    }
}
```

先程の例では、まず`Add`の`BinOp`が`ast`として入力され、`rhs`を評価する際に`Mul`の`BinOp`が再帰的に`ast`として入力されることになります。`Mul`の`BinOp`の計算結果を`lhs`である`1`と足し合わせて計算完了です。このようにして最終的に出てきた値が、入力した式の答えになります。


# JITコンパイラによる実装

ではJITコンパイルです。環境としてはLinux/x86_64を前提とします。今回はこのような手順でJITコンパイルに対応してみようと思います。

1. コンパイルしたネイティブコードを書き込む領域を用意する。
2. コンパイルしたネイティブコードを書き込む。
3. ネイティブコードの領域を関数として実行する。

ここで、コンパイル処理の入力となるのはASTです。なので、AST生成までのコードはインタプリタによる実装と変わらないことになります。


### 書き込む領域の用意

文章で見てもいまいち分かりづらいと思うので、ここからはソースコードを見ていきます。まずは領域を用意するところです。

```rust
extern crate libc;
use libc::{c_void, c_int, size_t, PROT_READ, PROT_WRITE, PROT_EXEC};
use std::alloc::{alloc, dealloc, Layout};

extern "C" {
    fn mprotect(addr: *const c_void, len: size_t, prot: c_int) -> c_int;
}
...
struct Compiler {
    p_start: *mut u8,
    p_current: *mut u8,
}

const CODE_AREA_SIZE: usize = 1024;
const PAGE_SIZE: usize = 4096;

impl Compiler {
    unsafe fn new() -> Self {
        let layout = Layout::from_size_align(CODE_AREA_SIZE, PAGE_SIZE).unwrap();
        let p_start = alloc(layout);
        let r = mprotect(p_start as *const c_void, CODE_AREA_SIZE, PROT_READ|PROT_WRITE|PROT_EXEC);
        assert!(r == 0);
        Compiler {
            p_start,
            p_current: p_start,
        }
    }
...
}
```

`mprotect`というCの関数を呼びたいので、libcクレートをインポートして関数の宣言をしておきます。この関数はLinux環境での標準Cライブラリに含まれるシステムコールラッパーです。Rustは（Linux環境しか見てませんが）デフォルトで標準Cライブラリ（libc.so）をリンクして動くので、この関数を使うのにリンカへの指示は必要ありません。

通常メモリは安全性のために保護されています。保護というのはできることに制限があるということで、ファイルと同じように読み取り、書き込み、実行の3種類の権限をそれぞれ設定できます。通常であれば獲得したメモリ領域に実行可能フラグは付いていないため、領域内のネイティブコードを実行するには権限を変更しておく必要があります。それを可能にするのが`mprotect`というLinuxシステムコールです。システムコールの詳細は[mprotectのマニュアル](https://linuxjm.osdn.jp/html/LDP_man-pages/man2/mprotect.2.html)をご参照下さい。

さて、それでは関数本体を見ていきます。`Layout::from_size_align`で獲得したいメモリのレイアウトを定義し、`alloc`を使ってメモリを獲得します。領域のサイズは適当です。ページサイズについてはいろいろ細かい話もありますが、Linuxシステムなら4KiBにしておけば大体は大丈夫だと思います。メモリが獲得できたら`mprotect`を使って実行権限を与えます。


### ネイティブコードの書き込み

領域が用意できたので、ネイティブコードを生成していきます。構造的には先ほどのインタプリタで式を評価する関数と変わりません。

```rust
impl Compiler {
...
    unsafe fn gen_code_ast(&mut self, ast: NodeKind) {
        match ast {
            Number(n) => {
                self.push_code(&[0x6a, n]); // push {}
            },
            BinOp { kind, lhs, rhs } => {
                self.gen_code_ast(*lhs);
                self.gen_code_ast(*rhs);
                self.push_code(&[0x5f]); // pop rdi
                self.push_code(&[0x58]); // pop rax
                match kind {
                    Add => {
                        self.push_code(&[0x48, 0x01, 0xf8]); // add rax, rdi
                    },
                    Sub => {
                        self.push_code(&[0x48, 0x29, 0xf8]); // sud rax, rdi
                    },
                    Mul => {
                        self.push_code(&[0x48, 0x0f, 0xaf, 0xc7]); // imul rax, rdi
                    },
                    Div => {
                        self.push_code(&[0x48, 0x99]); // cqo
                        self.push_code(&[0x48, 0xf7, 0xff]); // idiv rdi
                    },
                }
                self.push_code(&[0x50]); // push rax
            },
        }
    }
}
```

`push_code`に与えている16進数の配列がネイティブコードになります。ここではx86_64の機械語です。一応対応するアセンブリ表現も横にコメントで付けてあります。（そうしないと自分でも分からなくなる。）アセンブリはIntel記法を採用しており、これ以降の説明も同様です。

ネイティブコードを実行するときにASTを解析して再帰的に、みたいなやり方は難しいので、ここではスタックマシンと呼ばれる方法でシーケンシャルに実行できる命令を生成していきます。今回作る電卓には2項演算子しか出てこないため、この演算を「スタックから2つ数値をPOPして答えをスタックにPUSHする」という動作で実装すればよさそうです。

再度`1 + 2 * 3`の例を考えてみます。アセンブリで表現すると以下のようなコードが生成されます。

```
push 1
push 2
push 3
pop rdi
pop rax
imul rax, rdi
push rax
pop rdi
pop rax
add rax, rdi
push rax
```

アセンブリの細かい解説をここではできませんが、そこまで難しい知識は必要ありません。`push`と`pop`はそのままスタック操作で、`rax`と`rdi`はレジスタの名前（変数名のようなもの）です。少し注意として、`add`や`imul`などの計算は2つのレジスタを取り、1つ目のレジスタに結果が格納されます。つまり`add rax, rdi`では、`rax`の中身と`rdi`の中身を足し合わせた結果が`rax`に格納されます。


このコードが動く様子をスタックの状態で見てみます。まずは先頭の`1`がスタックにPUSHされます。

|&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;|
|:-:|
| 1 |

この1はそのままで`2 * 3`の計算が先に実行されます。左辺の`2`と右辺の`3`がそれぞれPUSHされます。

|&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;|
|:-:|
| 3 |
| 2 |
| 1 |

ここで2項演算子の乗算が行われます。スタックの頭にある2つの数値`3`と`2`をPOPし、乗算の結果である`6`がスタックにPUSHされます。

|&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;|
|:-:|
| 6 |
| 1 |

続いて2項演算子の加算が行われます。同様にスタックから`6`と`1`をPOPし、結果である`7`がPUSHされます。

|&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;|
|:-:|
| 7 |

ちなみに`push_code`はこのように実装しています。ループとかキャストとか、あまりいい実装じゃない気もしますね。

```rust
impl Compiler {
...
    unsafe fn push_code(&mut self, code: &[u8]) {
        for b in code.iter() {
            std::ptr::write(self.p_current, *b);
            self.p_current = (self.p_current as u64 + 1) as *mut u8;
        }
    }
...
}
```

ここまでで計算自体はできていますが、出来上がったネイティブコードの最後に`pop rax`と`ret`を追加します。

```rust
impl Compiler {
...
    unsafe fn gen_code(&mut self, ast: NodeKind) {
        self.gen_code_ast(ast);
        self.push_code(&[0x58]); // pop rax
        self.push_code(&[0xc3]); // ret
    }
}
```

これは、生成したネイティブコードを関数として実行させるための処理です。計算を実行すると、最終的な結果はスタックに残った状態になるため、この最終結果を取り出して`rax`に格納します。`ret`は関数の呼び出し元に戻るための命令です。ここで重要なのは、関数を呼び出すときに呼び出し元は呼び出した関数の戻り値が`rax`に格納されていることを期待している、ということです。このような仕様は[ABI（Application Binary Interface）](https://ja.wikipedia.org/wiki/Application_Binary_Interface)と呼ばれ、CPU、OS、言語などによって違ってきます。今回はRustのABIに従うことで、自分が組み立てたネイティブコードをあたかもRustの関数であるかのように使うことができます。


### 関数として実行

最後に生成したネイティブコードをRustのコードから呼び出します。ネイティブコードが格納されている領域はバイナリレベルでは実行形式における関数定義と変わりませんが、言語レベルでは関数として見えていないため、うまく呼び出すことができません。コード領域のポインタを関数ポインタに型変換する必要があります。このような型変換は、C言語などでは普通にキャストすればできますが、Rustでは`*mut u8`から`fn() -> u8`への`as`によるキャストが認められていません。このような場合には、標準ライブラリにある`transmute`という任意の型変換を行うことができる関数が使えます。あとは型変換した関数を実行し、戻り値を受け取れば計算結果が格納されているはずです。

```rust
pub fn interpret(line: &str, use_jit: bool) -> Result<u8, ()> {
...
        unsafe {
            let mut compiler = Compiler::new();
            compiler.gen_code(*ast);
            let code: fn() -> u8 = std::mem::transmute(compiler.p_start);

            // Run generated code!
            let r = code();

            compiler.free();
            Ok(r)
        }
...
}
```

今回、関数の型を`fn() -> u8`としましたが、これはABIとして「RustのABI」を使うことを意味しています。Rustでは関数のABIを[externにより指定](https://doc.rust-lang.org/reference/items/external-blocks.html)可能です。今回はシンプルな関数なのでそこまで意識せずに動作しましたが、もう少し複雑なことをする場合は、RustのABIとして正しく構築できているか、異なるABIを使うなら適切なexternを指定できているか、などABIへの意識が必要になると思います。


# ソースコード

ソースコード全体は[こちら](https://github.com/0yoyoyo/calculator)に上げていますので、よかったら見てみて下さい。


# 感想

JITコンパイルを理解するために実装してみましたが、書いているといろいろと疑問が出てきて、既存のJITコンパイラの実装などもいろいろ見てみたくなりました。あと、冒頭にも書いた通り電卓程度を実装してもあんまり効果が分からないので、もう少し頑張って性能検証とかができるレベルまでもっていきたいと思いました。