バイトコードコンパイラを実装する——スタックスロットとバックパッチングでASTを命令列に変換する

ASTをそのまま実行しない理由について

これまでの実装で、ソースコードは字句解析・構文解析を経てAST(Stmt/Exprの木構造)になり、型検査でその妥当性を確認できるようになりました。この状態のASTをそのまま再帰的に辿って評価する「木を歩くインタプリタ(tree-walking interpreter)」を書けば、プログラムを実行すること自体は可能です。しかし木を歩く実行方式は、ノードを辿るたびに仮想関数呼び出しやパターンマッチのコストがかかり、同じ式を繰り返し評価するループ処理などで顕著に遅くなります。

そこで多くの実用的な言語処理系が採用しているのが、ASTを一度「バイトコード」と呼ばれる単純な命令列にコンパイルしてから、その命令列をスタックマシンで実行するという2段階の方式です。今回はこのうちの前半、ASTをバイトコードにコンパイルする部分を実装しました。命令列を実際に実行するスタックマシンVMはまだ実装しておらず、代わりにコンパイル結果を人間が読める形に変換する逆アセンブラ(disassembler)を用意し、それを使って正しいバイトコードが生成されていることを確認しています。

バイトコードの表現形式——バイト列と定数プール

オペコード(命令の種類)とオペランド(命令の引数)を、それぞれ独立したRustの値として保持するのではなく、JVMやCPython、Lua、そしてこのプロジェクトが設計の参考にしているclox(Crafting InterpretersのバイトコードVM)と同じように、単純なVec<u8>に詰め込む方式を選びました。

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum OpCode {
    Constant,
    True,
    False,
    Pop,
    GetLocal,
    SetLocal,
    Add,
    Sub,
    // ...(算術・比較・論理演算子が続く)
    Len,
    MakeArray,
    Index,
    Jump,
    JumpIfFalse,
    Loop,
    Return,
}

#[derive(Debug, Clone, Default)]
pub struct Chunk {
    pub code: Vec<u8>,
    pub constants: Vec<Value>,
    pub lines: Vec<usize>,
}

OpCodeの各バリアントは1バイトの命令としてcodeに書き込まれ、オペランドを持つ命令はその直後に1〜2バイトのオペランドが続きます。例えばOP_CONSTANTは「定数プール(constants)の何番目の値を積むか」を表す1バイトのインデックスを、ジャンプ系の命令は「何バイト先へジャンプするか」を表す2バイトのオフセットを、それぞれ直後に持ちます。

Rustらしく書くなら、オペランドを持たせたenum(Jump(usize)のような)のリストとして表現する方がずっと型安全ですが、あえてバイト列方式を選びました。実際のバイトコードVMがどのようにメモリ効率よく命令を表現するかという標準的な技術を体験しておきたかったことに加え、後々ネイティブコード出力を試す際にも「機械語に近い低レベル表現を組み立てる」という感覚がそのまま活きると考えたためです。ただしその代償として、オペランドをバイト列から読み戻す処理(後述の逆アセンブラ)を自分で書く必要があります。

整数・浮動小数点数・真偽値・文字列・配列というこの言語の実行時の値は、Valueという1つのenumにまとめました。

#[derive(Debug, Clone, PartialEq)]
pub enum Value {
    Int(i64),
    Float(f64),
    Bool(bool),
    Str(String),
    Array(Vec<Value>),
}

配列をVec<Value>として値そのものに埋め込んでいるのは、この言語にはまだ代入文(x = 2;)も配列要素への代入(a[0] = 1;)も無く、配列を値として複製しても困る場面がないためです。もし将来ミュータブルな配列を導入するなら、Rc等を使った参照的な表現への作り直しが必要になりますが、今の時点でそこまで作り込むのは時期尚早と判断しました。

逆アセンブラ——VMがなくてもコンパイル結果を確認する

バイトコードは人間にとってはただのバイト列なので、正しく生成できているかを検証する手段が必要です。VMの実装はまだ先の作業なので、今回はChunkにバイト列を人間が読めるニーモニックへ変換するdisassembleを実装しました。

pub fn disassemble(&self, name: &str) -> String {
    let mut out = format!("== {name} ==\n");
    let mut offset = 0;
    while offset < self.code.len() {
        offset = self.disassemble_instruction(&mut out, offset);
    }
    out
}

先頭のバイトをOpCodeに戻す処理は、コンパイラが生成したバイト列を前提にしている(=不正なバイト列が渡ってくることは起こり得ない)ため、Resultは使わずpanicで済ませています。

impl OpCode {
    fn decode(byte: u8) -> OpCode {
        match byte {
            byte if byte == OpCode::Constant as u8 => OpCode::Constant,
            byte if byte == OpCode::True as u8 => OpCode::True,
            // ...(残りの全バリアントも同様に列挙)
            other => panic!("invalid opcode byte: {other}"),
        }
    }
}

命令ごとにオペランドの有無・バイト数が異なるため、ニーモニックへの変換も「オペランドなし」「1バイトのローカル変数スロット」「1バイトの定数インデックス」「2バイトのジャンプオフセット」の4パターンに分けています。ジャンプ命令の逆アセンブルでは、オフセットから実際のジャンプ先アドレスを計算して表示するようにしました。

fn jump_instruction(&self, out: &mut String, name: &str, sign: i32, offset: usize) -> usize {
    use std::fmt::Write;
    let jump = u16::from_be_bytes([self.code[offset + 1], self.code[offset + 2]]) as i32;
    let target = offset as i32 + 3 + sign * jump;
    writeln!(out, "{name:<16} {offset:4} -> {target}").unwrap();
    offset + 3
}

sign1を渡せば前方ジャンプ(OP_JUMP/OP_JUMP_IF_FALSE)、-1を渡せば後方ジャンプ(OP_LOOP)のアドレス計算になります。

ローカル変数はすべてスタックスロットで解決する

まだ関数(fn)が実装されていないため、この言語には「グローバルスコープ」と「関数内のローカルスコープ」という区別自体が存在しません。トップレベルのletもif/whileブロックの中のletも、区別なく「コンパイル時に決まるスタック上の位置(スロット番号)」として扱うことにしました。

struct Local {
    name: String,
    depth: usize,
}

pub struct Compiler {
    chunk: Chunk,
    locals: Vec<Local>,
    scope_depth: usize,
    loops: Vec<LoopContext>,
}

let x = 1;をコンパイルすると、まず右辺の式1を評価するバイトコードが積まれ、その後localsxという名前を追記するだけです。何か特別な「変数を定義する」命令を発行するわけではありません。式を評価した結果がすでにスタックの正しい位置に積まれているので、それがそのままxの実体になります。

StmtKind::Let { name, value } => {
    self.compile_expr(value, line, column)?;
    self.declare_local(name.clone());
    Ok(())
}

fn declare_local(&mut self, name: String) {
    self.locals.push(Local { name, depth: self.scope_depth });
}

変数の参照は、localsを後ろから(=内側のスコープから)探して最初に見つかった名前のインデックスをスロット番号として使います。

fn resolve_local(&self, name: &str) -> Option<usize> {
    self.locals.iter().rposition(|local| local.name == name)
}

後ろから探すことで、シャドーイング(let x = 1; if (true) { let x = 2; x; }のように内側で同じ名前をletし直す)も特別な処理なしに解決できます。内側のxlocalsの末尾に追記されているので、rpositionは自然と内側のxを先に見つけます。

ブロックに入るときと出るときはbegin_scope/end_scopeで対称に管理し、スコープを抜けるタイミングでそのスコープに属するローカル変数の分だけOP_POPを発行してスタックから取り除きます。

fn end_scope(&mut self, line: usize) {
    self.scope_depth -= 1;
    while let Some(local) = self.locals.last() {
        if local.depth <= self.scope_depth {
            break;
        }
        self.locals.pop();
        self.chunk.write_op(OpCode::Pop, line);
    }
}

clox本来の設計では、トップレベルの変数だけは名前引きのハッシュテーブル(グローバル変数テーブル)で管理し、関数の中のローカル変数だけをスタックスロットにします。これはREPLで1行ずつ評価しても後から定義した変数を先に参照できるようにするための仕組みですが、この言語にはまだREPLがなく、トップレベルの変数もスクリプト全体の中で宣言順に解決できれば十分です。そのため今回はグローバル変数テーブルを作らず、全部をスタックスロット方式に統一しました。関数を導入する際にどのみち必要になる仕組みを、二重に作らずに済んでいます。

制御フロー——ジャンプ命令とバックパッチング

if文は、条件式の結果に応じてthen節をスキップするかどうかを決めるOP_JUMP_IF_FALSEと、then節の実行後にelse節をスキップするOP_JUMPの組み合わせで表現します。

StmtKind::If { condition, then_branch, else_branch } => {
    self.compile_expr(condition, line, column)?;
    let then_jump = self.emit_jump(OpCode::JumpIfFalse, line);
    self.chunk.write_op(OpCode::Pop, line);
    self.compile_block(then_branch, line)?;
    let else_jump = self.emit_jump(OpCode::Jump, line);

    self.patch_jump(then_jump);
    self.chunk.write_op(OpCode::Pop, line);
    if let Some(else_stmt) = else_branch {
        self.compile_stmt(else_stmt)?;
    }
    self.patch_jump(else_jump);
    Ok(())
}

ここで難しいのは、OP_JUMP_IF_FALSEを発行する時点では「条件が偽だったときにどこへ飛べばよいか」がまだ分からないという点です。then節の中身がどれだけのバイト数になるかは、実際にコンパイルしてみないと分かりません。そこで標準的に使われているのが、いったんダミーのオフセット(0xFFFF)を書いておき、後から本当のジャンプ先が分かった時点でそのバイト列を書き換える「バックパッチング(backpatching)」という手法です。

fn emit_jump(&mut self, op: OpCode, line: usize) -> usize {
    self.chunk.write_op(op, line);
    self.chunk.write_u16(0xFFFF, line);
    self.chunk.code.len() - 2
}

fn patch_jump(&mut self, offset: usize) {
    let jump = self.chunk.code.len() - offset - 2;
    let jump = u16::try_from(jump).expect("jump distance exceeds u16 range");
    self.chunk.code[offset] = (jump >> 8) as u8;
    self.chunk.code[offset + 1] = jump as u8;
}

emit_jumpはダミーのオフセットを書いた「オペランドの位置」を返し、patch_jumpはその位置を後から呼び出した時点のcodeの長さで書き換えます。while文の条件チェックへ戻る後方ジャンプ(OP_LOOP)も原理は同じですが、ジャンプ先がすでに分かっている(条件式の直前にすでにコンパイル済み)ので、バックパッチングは不要です。

fn emit_loop(&mut self, loop_start: usize, line: usize) {
    self.chunk.write_op(OpCode::Loop, line);
    let offset = self.chunk.code.len() - loop_start + 2;
    let offset = u16::try_from(offset).expect("loop body exceeds u16 range");
    self.chunk.write_u16(offset, line);
}

短絡評価——&&||

&&||は、両辺を評価してから論理演算を適用するのではなく、右辺の評価を条件付きでスキップする短絡評価として実装する必要があります。false && f()f()が呼ばれてしまっては困るからです(この言語にはまだ関数呼び出しの実行自体がありませんが、意味論としては短絡評価が標準です)。

fn compile_and(&mut self, left: &Expr, right: &Expr, line: usize, column: usize) -> Result<(), CompileError> {
    self.compile_expr(left, line, column)?;
    let end_jump = self.emit_jump(OpCode::JumpIfFalse, line);
    self.chunk.write_op(OpCode::Pop, line);
    self.compile_expr(right, line, column)?;
    self.patch_jump(end_jump);
    Ok(())
}

左辺が偽であれば、右辺を評価する部分をまるごとジャンプでスキップし、スタックに残っている左辺のfalseをそのまま式全体の結果として扱います。左辺が真であれば、その値をOP_POPで捨てて右辺の評価に進み、右辺の値を式全体の結果とします。||はこの逆で、左辺が真であれば右辺をスキップし、偽であれば右辺を評価します。

for-inループの脱糖——インデックスループへの展開

for (item in arr) { ... }という構文には対応する専用のバイトコード命令を用意せず、代わりに「配列・添字・配列長」という3つの隠しローカル変数を使ったインデックスベースのwhile相当のバイトコードへ展開(脱糖、desugaring)しています。

fn compile_for(&mut self, var: &str, iterable: &Expr, body: &Block, line: usize, column: usize) -> Result<(), CompileError> {
    self.begin_scope();

    self.compile_expr(iterable, line, column)?;
    self.declare_local("<for-array>".to_string());
    let array_slot = (self.locals.len() - 1) as u8;

    self.emit_constant(Value::Int(0), line);
    self.declare_local("<for-index>".to_string());
    let index_slot = (self.locals.len() - 1) as u8;

    self.emit_get_local(array_slot, line);
    self.chunk.write_op(OpCode::Len, line);
    self.declare_local("<for-length>".to_string());
    let length_slot = (self.locals.len() - 1) as u8;

    // ...(添字 < 配列長 の間ループし、本体の中で
    //      arr[添字] を `item` として束縛する)
}

配列の長さを取るOP_LENは、この脱糖のためだけに新設した命令です。ループ本体を実行した後は、添字用のローカル変数を1つ増やしてOP_LOOPで条件チェックへ戻ります。この増分処理にはOP_SET_LOCALを使っていますが、この言語にはまだx = 1;という代入構文自体が存在しません。ここでのOP_SET_LOCALは、表面上の文法とは独立にコンパイラが内部的に発行しているものです。

break/continueとスタックの整合性

型検査のフェーズでは先送りにしていた「break/continueがループの外にないか」の検査を、このコンパイラで実装しました。コンパイラはループに入るたびにLoopContextをスタックに積んでおき、break/continueが出てきた時点でこのスタックが空であればコンパイルエラーにします。

StmtKind::Break => {
    let Some(loop_ctx) = self.loops.last() else {
        return Err(self.error(line, column, "'break' outside of a loop".to_string()));
    };
    // ...
}

もう一つ注意が必要なのが、break/continueによってブロックの外へジャンプする際、そのブロックの中でletされたローカル変数がスタックに残ったままにならないようにすることです。通常、スコープを抜けるときはend_scopeが対応するOP_POPを発行しますが、break/continueはスコープの終わりを待たずに途中でジャンプしてしまうため、end_scopeによるPOPを素通りしてしまいます。そこで、ループに入った時点でのローカル変数の数をLoopContextに記録しておき、break/continueの時点でそれより後に増えた分だけ明示的にOP_POPを発行してからジャンプするようにしました。

struct LoopContext {
    continue_target: ContinueTarget,
    break_jumps: Vec<usize>,
    locals_at_entry: usize,
}

fn emit_pop_locals_above(&mut self, count: usize, line: usize) {
    for _ in count..self.locals.len() {
        self.chunk.write_op(OpCode::Pop, line);
    }
}

continueのジャンプ先も一様ではありません。whileループでは条件式の直前(loop_start)へ戻ればよいので後方ジャンプで済みますが、forループでは添字を1つ進める増分処理を飛ばすわけにはいかないため、増分処理の開始位置(本体の直後)へ前方ジャンプする必要があります。増分処理はまだコンパイルされていない未来の位置なので、breakと同様にバックパッチング対象として保留しておき、増分処理をコンパイルし終えた時点でまとめて解決します。

enum ContinueTarget {
    Loop(usize),
    Forward(Vec<usize>),
}

動作確認

if/elseと、内部でletしたローカル変数を含むスクリプトをコンパイルし、逆アセンブル結果を出力してみます。

let x = 1;
let y = 2;
if (x < y) {
    let z = x + y;
    z;
} else {
    0;
}
$ cargo run -- examples/if_else.na
== script ==
0000    1 OP_CONSTANT         0 '1'
0002    2 OP_CONSTANT         1 '2'
0004    3 OP_GET_LOCAL        0
0006    | OP_GET_LOCAL        1
0008    | OP_LT
0009    | OP_JUMP_IF_FALSE    9 -> 25
0012    | OP_POP
0013    4 OP_GET_LOCAL        0
0015    | OP_GET_LOCAL        1
0017    | OP_ADD
0018    5 OP_GET_LOCAL        2
0020    | OP_POP
0021    3 OP_POP
0022    | OP_JUMP            22 -> 29
0025    | OP_POP
0026    7 OP_CONSTANT         2 '0'
0028    | OP_POP
0029    3 OP_RETURN

xはスロット0、yはスロット1として解決され、then節の中でletしたzはスロット2になっています。オフセット21のOP_POPは、then節を抜ける際にzをスタックから取り除くためのものです。オフセット9のOP_JUMP_IF_FALSEは条件が偽の場合にオフセット25(else節の直前)へ、オフセット22のOP_JUMPthen節の実行後にelse節をスキップしてオフセット29(OP_RETURN)へ、それぞれ正しくジャンプしています。

whileループの中でbreakする例も確認します。

let i = 0;
while (i < 10) {
    if (i == 3) {
        break;
    }
    i;
}
$ cargo run -- examples/while_break.na
== script ==
0000    1 OP_CONSTANT         0 '0'
0002    2 OP_GET_LOCAL        0
0004    | OP_CONSTANT         1 '10'
0006    | OP_LT
0007    | OP_JUMP_IF_FALSE    7 -> 33
0010    | OP_POP
0011    3 OP_GET_LOCAL        0
0013    | OP_CONSTANT         2 '3'
0015    | OP_EQ
0016    | OP_JUMP_IF_FALSE   16 -> 26
0019    | OP_POP
0020    4 OP_JUMP            20 -> 34
0023    3 OP_JUMP            23 -> 27
0026    | OP_POP
0027    6 OP_GET_LOCAL        0
0029    | OP_POP
0030    2 OP_LOOP            30 -> 2
0033    | OP_POP
0034    | OP_RETURN

break(オフセット20のOP_JUMP)と、ループが自然に終了する経路(オフセット7のOP_JUMP_IF_FALSEが偽になった場合)が、どちらも同じオフセット34(OP_RETURNの直前)に合流していることが確認できます。breakによるジャンプが、通常の終了経路と同じ場所へ着地するように設計した通りの結果です。

最後に、forループがインデックスループへ展開されている様子です。

for (n in [1, 2, 3]) {
    n;
}
$ cargo run -- examples/for_loop.na
== script ==
0000    1 OP_CONSTANT         0 '1'
0002    | OP_CONSTANT         1 '2'
0004    | OP_CONSTANT         2 '3'
0006    | OP_MAKE_ARRAY       3
0008    | OP_CONSTANT         3 '0'
0010    | OP_GET_LOCAL        0
0012    | OP_LEN
0013    | OP_GET_LOCAL        1
0015    | OP_GET_LOCAL        2
0017    | OP_LT
0018    | OP_JUMP_IF_FALSE   18 -> 42
0021    | OP_POP
0022    | OP_GET_LOCAL        0
0024    | OP_GET_LOCAL        1
0026    | OP_INDEX
0027    2 OP_GET_LOCAL        3
0029    | OP_POP
0030    1 OP_POP
0031    | OP_GET_LOCAL        1
0033    | OP_CONSTANT         4 '1'
0035    | OP_ADD
0036    | OP_SET_LOCAL        1
0038    | OP_POP
0039    | OP_LOOP            39 -> 13
0042    | OP_POP
0043    | OP_POP
0044    | OP_POP
0045    | OP_POP
0046    | OP_RETURN

配列リテラル(オフセット0〜7)に続いて添字用の0(オフセット8)と配列長(オフセット10〜12)を積み、オフセット13から始まる条件チェックと本体、オフセット31〜38の増分処理、そしてオフセット39のOP_LOOPでオフセット13へ戻るという、狙い通りの構造になっています。

コンパイラ自体のテストは58件(既存のLexer/Parser/型検査のテストを含む)が通ることを確認しています。

備考

定数プールのインデックスは1バイト(最大256個)

OP_CONSTANTのオペランドを1バイトにしているため、1つのスクリプトに書ける定数リテラルは256個までという制約があります。clox(Crafting Interpreters)では定数が256個を超えた場合に3バイトオペランドのOP_CONSTANT_LONGへ切り替える仕組みが用意されていますが、今の学習用スクリプトの規模でこの上限に達することはまず無いと判断し、今回は見送りました。実際に上限を超えるとadd_constant内のexpectでpanicします。

returnはまだ「値を返す」という意味を持たない

関数(呼び出し元となるコールスタック)がまだ存在しないため、returnは今のところ実質的に「そこでコンパイルを打ち切り、OP_RETURNを発行するだけ」の命令として扱っています。return式;の場合、式は評価してからOP_POPで捨てており、呼び出し元に値を持ち帰る仕組みはありません。関数を実装する段階で、コールフレームに値を返す本来の意味へ作り直す予定です。

cargo buildValue::BoolValue::Arrayが未使用という警告が出る

warning: variants `Bool` and `Array` are never constructed

真偽値リテラルはOP_TRUE/OP_FALSEという専用の命令にしており、配列リテラルもOP_MAKE_ARRAYで実行時に組み立てる設計にしているため、コンパイラのコードの中でValue::BoolValue::Arrayを直接構築する箇所がありません。これらのバリアントは、実行系(VM)がOP_TRUE/OP_MAKE_ARRAYを処理してスタックに積む段階で初めて使われるようになります。VMが実装されるまでの間、この警告は無視して問題ありません。