🌆

Rustによる言語内x86_64アセンブラのAPIデザインと実装

2022/11/20に公開

Rust

assembly

tech

概要

Rustでx86_64アセンブラを実装しました。JITコンパイラの実装での利用を想定しているため、アセンブリコードを文字列として入力に取るのではなく、 静的に型付けされたRustプログラム上にアセンブリコードを記述できるライブラリ として実装しています。例えばフィボナッチ数を計算するアセンブリコードを以下のように記述し、付属のJITエンジンを用いて実行できます:

fn fib_object() -> io::Result<Object> {
    let mut w = Writer::new();
    let fib = w.get_label("fib");
    let l1 = w.issue_label();
    let l2 = w.issue_label();

    w.define(fib, true);
    w.cmpl(Edi, 1i8)?;
    w.jle(Short(l2))?;
    w.movl(Edx, 1)?;
    w.movl(Eax, 1)?;
    w.xorl(Ecx, Ecx)?;

    w.define(l1, false);
    w.movl(Esi, Eax)?;
    w.addl(Edx, 1i8)?;
    w.addl(Eax, Ecx)?;
    w.movl(Ecx, Esi)?;
    w.cmpl(Edi, Edx)?;
    w.jne(Short(l1))?;
    w.retq()?;

    w.define(l2, false);
    w.movl(Eax, 1)?;
    w.retq()?;

    w.produce()
}

#[test]
fn fib_jit() {
    let obj = fib_object().unwrap();

    let mut engine = jit::Engine::new(symbol_resolver::none);
    assert!(engine.add_object(&obj).is_ok());

    let fib = engine.get("fib").expect("fib not defined");
    let fib = unsafe { std::mem::transmute::<_, extern "C" fn(i32) -> i32>(fib) };
    assert_eq!(fib(10), 55);
}

本記事では、このライブラリのRust上でのAPIデザインや実装について振り返りつつ解説します。

導入

※本筋ではないため読み飛ばしても問題ありません。

動機

アセンブラはよく枯れた実装がいくつか存在しますが、今回は以下のような動機から新たに作成することにしました。

(x86_64 ISAの) オペランドで表現できるアドレッシング等が複雑で、どのようにエンコードされるか知りたくなった
binutilsでは抽象化されているように見えたcallやjmpなどにおけるアドレッシングや、ラベルとリンカとの関係などを実装して理解したかった
アセンブラは自作言語処理系での利用を目的としていたが、自作言語には言語機能としてLisp-likeなマクロがあり、これを実現するためのJITコンパイルに際して、プロセス外のアセンブラ(as)を呼び出して生成結果のオブジェクトをメモリにロードして…といった部分をプロセス内で完結するようにしてみたかった

学習が主な動機となっています。一方で、3つ目の理由から、テキストファイルの入力サポートなどは持たない、Rustコード上で記述できるライブラリとしてアセンブラを実装することにしました。

アセンブリ言語とエンコーディングを理解する

もともと自分はアセンブリコードをほとんど書いたことがなく、アセンブリ言語への馴染みが薄い状態でした。また、アセンブリの疑似命令が最終的な出力にはどのように反映されるか、リンカとどのように協調しているかなどの理解がぼんやりしていると感じていました。そのため、まずは以下のような資料を漁り、アセンブリコードと機械語のギャップや、エンコーディング、リンク時の動作を把握していきました。

簡単な関数をC言語で書いて objdump -Str -M intel してみる
既存のアセンブラ実装を見る
- 特にStanfordPL/x64asmが小さくて参考になりました
OSDev WikiのX86-64 Instruction Encodingの記事を読む
- 最初は全くわかりませんでしたが、これはx86_64のエンコードが歴史的な事情もあり非常に複雑なためで、一つずつ噛み砕いていくことで理解できるようになりました
Intel SDMを読む
chibiccのソースコードリーディング

このフェーズは本記事の主題ではないですが、実装を含めても このフェーズが一番困難な工程でした。

APIデザインを考える

Rustコード上でアセンブリ言語を表現する上で、ライブラリが提供するAPIにはいくつかの形式が考えられます。

let foo = ...; // 埋め込む即値など

// 1. 文字列を入力に取る関数を提供する
assemble(format!(r##"
  movl esi, eax
  addl eax, {}
"##, foo))

// 2. アセンブリコードを記述できるマクロを提供する
asm! {
  movl esi, eax
  addl eax, {foo}
}

// 3. Rustプログラムとして記述できるようなAPIを提供する
asm(|w| {
  w.movl(Esi, Eax);
  w.addl(Eax, foo);
})

1は as コマンドを呼び出すのとそう変わりありません。

2はコンパイル時に (また、IDE上でリアルタイムに) アセンブリコードが適切であるか検証できます。また、 (埋め込みによってランタイムに結果が変わる部分を除いて) エンコードまでコンパイルタイムに済ませることで、ランタイムのオーバーヘッドを最小にすることもできそうです。ただ、この形式はアセンブリのパース処理などがあり実装は大変です。

今回は3のような形式を採用することにしました。ニーモニックはメソッドとして、オペランドはメソッドの引数として表現します。単なるRustプログラムなのでIDEの補完がよく効き、不適切なオペランド指定なども型エラーとしてIDE上ですぐ確認できます。記述が若干冗長ですが、今回の目的…言語処理系のバックエンド実装では、長大なアセンブリコードを扱うわけではなく、小さなアセンブリコード片がたくさん存在するのみというのも理由に挙げられます。

impl Codegen {
    ...
    // 言語処理系実装の疑似コード: 加算演算子の実装
    fn eval_binary(&mut self, op: &Binary, a: &Expr, b: &Expr) -> io::Result<Layout> {
        use Binary::*;

        let a_ret = self.eval(a)?; // left operand
        self.push(&a_ret)?;
        let b_ret = self.eval(b)?; // right operand

        match op {
            SAdd | UAdd | SSub | USub | ... => {
                assert_eq!(a_ret.size, b_ret.size);
                self.w.movq(Rcx, Rax)?; // right operand
                self.pop(&a_ret)?;   // left operand

                match (op, b_ret.size) {
                    (SAdd | UAdd, 1) => self.w.addb(Al, Cl)?,
                    (SAdd | UAdd, 2) => self.w.addw(Ax, Cx)?,
                    (SAdd | UAdd, 4) => self.w.addl(Eax, Ecx)?,
                    (SAdd | UAdd, 8) => self.w.addq(Rax, Rcx)?,
                    ...
                }
            }
            ...
        }
    }
    ...
}

ということで、このようなニーモニック (addl) やオペランド (Eax, Rcx など) をそれぞれ実装していくことになります。

実装

オペランドをRustで表現する

src/asm/operand.rs

x86_64 ISAでは、命令が取れるオペランドにはいくつかの種類があります。例としてadd(加算)命令を見てみます。

Instruction	Encoding	Description
`addl r32, r/m32`	`03 /r`	Add r/m32 to r32.
`addq r64, r/m64`	`REX.W+ 03 /r`	Add r/m64 to r64.
`addl r/m32, r32`	`01 /r`	Add r32 to r/m32.
`addq r/m64, r64`	`REX.W+ 01 /r`	Add r64 to r/m64.
`addl r/m32, imm32`	`81 /0 id`	Add imm32 to r/m32.
`addq r/m64, imm32`	`REX.W+ 81 /0 id`	Add imm32 sign-extended to 64-bits to r/m64.

Instruction中の r32 や r/m32 のような構文がオペランドに取るものを表現しています。

r32 はdoublewordな汎用レジスタ (eax, edx, esi, r10d など)
r64 はquadwordな汎用レジスタ (rax, rdx, rdi, r10 など)
r/m64 はquadwordな汎用レジスタ、またはメモリオペランド ([rdi], [rdi + rsi] など; 後述)
imm32 は32bitの即値 (123, -4567 など)

まずはこれらをRust上で表現していきます。

汎用レジスタは単純で、種類ごとにenumで定義します。例えば r64 なら:

// quadwordな汎用レジスタ (64bit general-purpose register)
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
pub enum Gpr64 {
    Rax,
    Rcx,
    Rdx,
    ...
}
pub use Gpr64::*;

即値 imm32 はRustの整数型をそのまま用います。

レジスタまたはメモリオペランドを取る r/m64 については、それをそのまま表現する型を設ける代わりに、オペランドの種類ごとに定義を展開します。例えば addq r/m64, imm32 なら、 addq r64, imm32 と addq m64, imm32 の2つ定義があると考えます。

メモリオペランドをRustで表現する

多くのオペランドは上述のようにシンプルに表現できますが、メモリオペランドはもう少し複雑です。x86_64のメモリオペランドは表現力が高く、例えば以下のような表現が可能です。

[disp32] (disp32は32bit整数)
[RIP + disp32] (RIP相対)
[disp32 + r64 * scale] (r64は汎用レジスタ、scaleは1, 2, 4, または8)
[r64 + r64 * scale]
[r64 + disp32 + r64 * scale]

最後の例が表現可能な最も複雑な形で、アドレスが「ベースレジスタ + ディスプレースメント値 + インデックスレジスタ * スケール値」で計算されます。このようなアドレスを柔軟に表現できるよう、Rustでは型パラメタ付きのstructで表現することにします。

/// メモリアドレスの表現。 `Base + Disp + Idxs`
/// 指定がない部分の型パラメタは () で埋める
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy, new)]
pub struct Address<Base, Disp, Idxs> {
    pub base: Base,
    pub disp: Disp,
    pub idxs: Idxs,
}

/// `r64 * scale` の表現
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
pub struct IndexScale {
    pub index: Gpr64,
    pub scale: Scale,
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
pub enum Scale {
    Scale1,
    Scale2,
    Scale4,
    Scale8,
}

// 以下のような記述・型付けを目指す:
let a = Rax - 32;          // OK! (a: Address<Gpr64, i32, ()>)
let b = Rax + Rcx * 4;     // OK! (b: Address<Gpr64, (),  IndexScale>)
let c = Rax + Rcx + Rdx;   // type error!
let d = Rax * 2 + Rcx * 2; // type error!
let e = Rax + Rcx + 123;   // OK! (e: Address<Gpr64, i32, IndexScale>)

Rustは std::ops::Add などのtraitを実装することで演算子のオーバーロードが可能なので、このような記述が可能なAPIを提供することができます。

// r64 * u8
impl Mul<u8> for Gpr64 { type Output = IndexScale; ...  }
impl Mul<Scale> for Gpr64 { type Output = IndexScale; ...  }

// r64 + i32, r64 - i32
impl Add<i32> for Gpr64 { type Output = Address<Gpr64, i32, ()>; ... }
impl Sub<i32> for Gpr64 { type Output = Address<Gpr64, i32, ()>; ... }

// r64 + r64 * scale
impl Add<Gpr64> for Address<Gpr64, (), ()> {
  type Output = Address<Gpr64, (), IndexScale>;
  ...
}

// r64 + i32 + r64 * scale
impl Add<Gpr64> for Address<Gpr64, i32, ()> {
  type Output = Address<Gpr64, i32, IndexScale>;
  ...
}
... // 似たような定義が続く

入力に対して、出力の type Output では Address の埋まっている型パラメタが増えていることが読み取れると思います。

最後に、アドレスからメモリオペランドを構築できるようにします。

/// メモリオペランド。Intel構文の `[..]` と同等
/// 内部的にはエンコード済みの (そしてパラメタ化された型をeraseした) アドレスを持つ
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
pub struct Memory(pub EncodedAddress);

pub fn memory(addr: impl Into<EncodedAddress>) -> Memory {
    Memory(addr.into())
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
pub struct EncodedAddress(..);

// メモリアドレスとして適切な構文に対応する型 `T` について `From<T> for EncodedAddress` を実装する
impl From<i32> for EncodedAddress { ... }
impl From<Gpr64> for EncodedAddress { ... }
impl From<Address<(), i32, ()>> for EncodedAddress { ... }
impl From<Address<(), i32, IndexScale>> for EncodedAddress { ... }
impl From<Address<Gpr64, i32, ()>> for EncodedAddress { ... }
...

これでオペランドの表現は一通り揃いました。ここから、これらオペランドを引数に取るメソッド(ニーモニック)を実装していきます。

ニーモニックとメソッドの対応: traitによるオーバーロード

src/asm/inst.rs

addq を例に取ります。 addq は以下のようなオペランドを取ることができます。

w.addq(Rax, Rcx)?;               // (Gpr64, Gpr64)
w.addq(memory(Rax + Rcx), Rdx)?; // (Memory, Gpr64)
w.addq(memory(Rdx), 123)?;       // (Memory, i32)

Rustにはいわゆる関数のオーバーロードはないので、traitで「 addq が (Gpr64, Gpr64), (Memory, Gpr64), ... について定義されている」というような表現を設ける必要があります。

/// 特定の命令をエンコードして `w` に書き込むtrait
pub trait WriteInst<W: ?Sized> {
    fn write_inst(&self, w: &mut W) -> io::Result<()>;
}

/// addq op1, op2 を Addq(op1, op2) で表現する
pub struct Addq<Op0, Op1>(pub Op0, pub Op1);

/// addq r64 r64: Add r/m64 to r64.
impl<W: io::Write + ?Sized> WriteInst<W> for Addq<Gpr64, Gpr64> { ... }
/// addq m64 r64: Add r64 to r/m64.
impl<W: io::Write + ?Sized> WriteInst<W> for Addq<Memory, Gpr64> { ... }
/// addq r64 imm32: Add imm32 sign-extended to 64-bits to r/m64.
impl<W: io::Write + ?Sized> WriteInst<W> for Addq<Gpr64, i32> { ... }
...

適切なオペランドの組からなる Addq<Op1, Op2> は WriteInst を実装している、という形になっています。次に、byteorder::WriteBytesExtのように io::Write を拡張する形でこれを呼び出せるようにします。

pub trait WriteInstExt: io::Write {
    ...
    fn addq<Op0, Op1>(&mut self, op0: Op0, op1: Op1) -> io::Result<()>
    where
        Addq<Op0, Op1>: WriteInst<Self>,
    {
        Addq(op0, op1).write_inst(self)
    }
    ...
}

impl<W: io::Write + ?Sized> WriteInstExt for W {}

これにより、 io::Write 上で以下のようにアセンブリを記述できるようになりました。

fn hello_asm(w: &mut (impl io::Write + ?Sized)) -> io::Result<()> {
    w.movq(Rcx, memory(Rax))?;         // OK
    w.movq(memory(Rdx), memory(Rcx))?; // type error!
    Ok(())
}

最後に

アセンブラを題材としつつ、RustのtraitによるAPIデザインの可能性を探る話となりました。本ライブラリの実装では、他にも以下のようなことを行っているので、アセンブラ関連の実装に興味がある方は参照していただければと思います。

ラベルの定義と解決
- Writer::produce 時に io::Seek を用いてラベル位置の書き込みを行う
- WriteInst を実装することで Writer 上で w.callq(Label) のような仮想命令を呼び出せるようにする
ELF形式の出力のサポート
JIT実行のサポート
テンプレートエンジンAskamaを用いた命令セットのコード生成