XbyakをRustから使ってみた

はじめに

XbyakはC++で書かれたJITアセンブラです。Xbyak自体は「そういうアセンブラがあるらしい」程度にしか知らなかったのですが、"JITアセンブラXbyakを使ってみる（その１）"を読んで「JITである（つまりプログラム実行時に動的にアセンブリを生成・実行する）」というのを知り、面白そうなので使ってみようと思いました。
とはいえここ数年Rustばかり書いていて、C++を書くのは億劫になってしまったためRustから使う方法を考えてみました。
本当はいい感じのインターフェースができればライブラリにして公開するつもりだったのですが、いまいちいい感じのが思いつかなかったため、作業履歴だけ残しておきます。どなたかが良いインターフェースを思いついてくれる（あわよくばコードも書いてくれる）ことに期待します。

方針

XbyakのコードはC++なので、別途.cppファイルに書いてC ABIを通してリンクすればいいのですが、やはりインラインアセンブラ的に書きたいです。したがって手続きマクロを使ってRustのコード中にXbyakのC++コードを埋め込みます。
すなわち

use xbyak_rs::xbyak;

xbyak! {
    struct Code : Xbyak::CodeGenerator {
        Code(int x)
        {
            mov(eax, x);
            ret();
        }
    };

    extern "C" int run(int n) {
        Code c(n);
        int (*f)() = c.getCode<int (*)()>();
        return f();
    }
}

のようなマクロxbyak!を用意して、このコードをコンパイル・リンクすればよいということになります。
（このマクロの中身はhttps://github.com/herumi/xbyak#how-to-use-itのサンプルコードと同じなので、コードの説明はそちらをご参照ください。）

手続きマクロ

xbyak!マクロの実装は以下のような感じになります。

use proc_macro::TokenStream;
use std::env;
use std::fs;
use std::path::Path;
use std::str::FromStr;

static XBYAK_H: &'static str = include_str!("xbyak/xbyak.h");
static XBYAK_BIN2HEX_H: &'static str = include_str!("./xbyak/xbyak_bin2hex.h");
static XBYAK_MNEMONIC_H: &'static str = include_str!("./xbyak/xbyak_mnemonic.h");
static XBYAK_UTIL_H: &'static str = include_str!("./xbyak/xbyak_util.h");
static XBYAK_COPYRIGHT: &'static str = include_str!("./xbyak/COPYRIGHT");

# [proc_macro]
pub fn xbyak(input: TokenStream) -> TokenStream {
    let out_dir = env::var_os("OUT_DIR").unwrap();

    let xbyak_path = Path::new(&out_dir).join("xbyak.h");
    fs::write(&xbyak_path, XBYAK_H).unwrap();

    let xbyak_path = Path::new(&out_dir).join("xbyak_bin2hex.h");
    fs::write(&xbyak_path, XBYAK_BIN2HEX_H).unwrap();

    let xbyak_path = Path::new(&out_dir).join("xbyak_mnemonic.h");
    fs::write(&xbyak_path, XBYAK_MNEMONIC_H).unwrap();

    let xbyak_path = Path::new(&out_dir).join("xbyak_util.h");
    fs::write(&xbyak_path, XBYAK_UTIL_H).unwrap();

    let xbyak_path = Path::new(&out_dir).join("COPYRIGHT");
    fs::write(&xbyak_path, XBYAK_COPYRIGHT).unwrap();

    let mut code = String::from("#include <xbyak.h>\n");
    code.push_str(&input.to_string());
    let dest_path = Path::new(&out_dir).join("xbyak_code.cpp");
    fs::write(&dest_path, code).unwrap();

    cc::Build::new()
        .cpp(true)
        .file(&dest_path)
        .include(&out_dir)
        .compile("xbyak_code");

    FromStr::from_str(&"").unwrap()
}

Xbyakの本体であるヘッダファイル群は手続きマクロにinclude_strで取り込み、手続きマクロの実行時にOUT_DIRに展開します。さらに、xbyak!マクロの中身を、先頭に#include <xbyak.h>を付与して書き出します。
書き出したxbyak_code.cppをccクレートを使ってコンパイルしています。

使い方

使う際にはまずbuild.rsが必要です。

build.rs

fn main() {
    println!(
        "cargo:rustc-env=OUT_DIR={}",
        std::env::var("OUT_DIR").unwrap()
    );
    println!(
        "cargo:rustc-env=TARGET={}",
        std::env::var("TARGET").unwrap()
    );
    println!(
        "cargo:rustc-env=OPT_LEVEL={}",
        std::env::var("OPT_LEVEL").unwrap()
    );
    println!("cargo:rustc-env=HOST={}", std::env::var("HOST").unwrap());

    println!("cargo:rustc-link-lib=static=xbyak_code");
    println!(
        "cargo:rustc-link-search=native={}",
        std::env::var("OUT_DIR").unwrap()
    );
    println!("cargo:rustc-link-lib=stdc++");
}

このあたりはライブラリ化するならxbyak_rs::build()のような感じでまとめる必要があるでしょう。
最初の4つのprintln!は手続きマクロ内でccによるコンパイルをするために使います。これらの環境変数はbuild.rsでのみ提供されており、手続きマクロからアクセスすることはできません。ccはbuild.rsで使われる想定なのでこれらの環境変数が必須であり、ここでrustcに渡しています。
後半の3つのprintln!は手続きマクロで生成したCのライブラリをリンクするためのオプションです。build.rsにてccを使ってコンパイルするとこのあたりは自動でやってくれるのですが、手続きマクロ内で使用した場合はそれが効かないので手動でやる必要があります。

これを用意すれば以下のようなコードが通るようになります。

use xbyak_rs::xbyak;

xbyak! {
    struct Code : Xbyak::CodeGenerator {
        Code(int x)
        {
            mov(eax, x);
            ret();
        }
    };

    extern "C" int run(int n) {
        Code c(n);
        int (*f)() = c.getCode<int (*)()>();
        return f();
    }
}

extern "C" {
    fn run(n: i32) -> i32;
}

fn main() {
    unsafe {
        println!("{}", run(45));
    }
}

Rust的なインターフェースを考える

とりあえず動作確認はできたのでRust的にどんな感じのインターフェースがいいのかを考えます。
少なくとも2か所のextern "C"は手続きマクロによる自動生成であって欲しいです。
コードジェネレータのオブジェクトCodeはRustでもstructで良さそうです。
すなわちこんなイメージでしょうか。

xbyak! {
    struct Code : Xbyak::CodeGenerator {
        Code(int x)
        {
            mov(eax, x);
            ret();
        }
    };
}

// xbyak!マクロによって以下が生成される
// struct Code {
//     fn new(i32) -> Self {
//         ...
//     }
//     fn run(&self) -> i32 {
//         ...
//     }
// }

fn main() {
    let code = Code::new(45);
    unsafe {
        println!("{}", code.run());
    }
}

問題はCodeのrunの引数と戻り値です。これはC++のコード上はc.getCode<int (*)()>();の型引数int (*)()に対応するもので、この場合は引数なしで戻り値がintになります。
なのでrunの関数プロトタイプをRustで宣言し、それを解析してC++の型引数を生成する感じになると思うのですが、どう宣言させるのがいいか…というところです。