ガベージコレクタを実装する——循環参照を作りうるクロージャとupvalueだけを生ポインタのヒープへ移し、mark-sweepで回収する

Rcでは解放できない参照が生まれる

関数とクロージャは、これまでRc(クロージャのupvalueはRc<RefCell<_>>)で持っていました。参照カウントによる管理は実装が単純で、これまでの言語機能の範囲では何の問題もありませんでした。しかしこの処理系は今後、構造体とメソッドを実装する予定です。メソッドがself(インスタンス自身)をupvalueとして捕まえるクロージャになり、そのクロージャをインスタンス側が保持する(バインドされたメソッドとしてキャッシュする、など)ようになると、次のような循環ができます。

// 構造体はまだ実装していません。将来的にこういう形になる予定、というイメージです。
struct Counter {
    count: int;
    fn increment(): void { self.count = self.count + 1; }
}

incrementが捕まえたselfがそのインスタンス自身を指し、インスタンス側もそのクロージャを(メソッドとして、あるいはフィールドとして)保持していれば、instance → closure → upvalue → instanceという循環参照ができあがります。Rcは参照カウントが0にならない限り解放しないので、この種の循環は永久にリークします。構造体を実装する前に、本物のmark-sweep GCへ土台を差し替えておく必要がありました。

対象を絞る——ヒープに移すのはClosureとUpvalueだけ

GCヒープに載せる対象は最小限に絞りました。Functionはコンパイル時に一度だけ作られる不変の値で、自分の持つバイトコードと定数プールの中にClosureが紛れ込むことは無いため、循環参照を作る経路がありません。一方ClosureUpvalueは実行時にしか作られず、両者の間で循環(クロージャが持つupvalueが、そのクロージャ自身を指す値になる)が起き得ます。

pub enum Value {
    Int(i64),
    Float(f64),
    Bool(bool),
    Str(String),
    Array(Vec<Value>),
    Function(Rc<FunctionObj>), // 循環参照を作れないのでRcのまま
    Closure(*mut Obj),         // GCヒープへのポインタに変更
    Void,
}

Functionを無改修のままにできたことで、GCヒープはVMの実行時にしか使われないことが確定しました。関数値はコンパイル時にチャンクの定数として埋め込まれる(CompilerRc::new(FunctionObj {...})を作る)のに対し、ClosureUpvalueOP_CLOSUREの実行時や、upvalueを捕まえる瞬間にVMが初めて生成します。この線引きのおかげで、コンパイラとVMでヒープを共有する必要が一切なくなり、Compiler側は無改修で済みました。

Objの表現——C言語の型パニングをRustのenumに置き換える

clox(Crafting Interpretersのバイトコード実装)は、あらゆるヒープオブジェクトの先頭に共通のObjヘッダを埋め込み、Obj*と各具象型(ObjString*など)のポインタを相互キャストして扱います。C言語では構造体の最初のフィールドが同じメモリレイアウトを持つことを利用したテクニックですが、Rustにこの型パニングを持ち込む理由はありません。代わりに、ペイロードをenumで持たせました。

// heap.rs
pub struct Obj {
    marked: bool,
    next: *mut Obj,
    pub data: ObjData,
}

pub enum ObjData {
    Closure(ClosureObj),
    Upvalue(UpvalueData),
}

pub struct ClosureObj {
    pub function: Rc<FunctionObj>,
    pub upvalues: Vec<*mut Obj>,
}

pub enum UpvalueData {
    Open(usize),
    Closed(Value),
}

nextが全オブジェクトを1本の侵入型単方向リストにつなぐ役割を持つのは、cloxのvm.objectsと同じ発想です。GCが「今ヒープにあるすべてのオブジェクト」を漏れなく辿るための唯一の手掛かりが、このリストです。

Objの確保はBox::into_rawで行い、確保のたびにこのリストの先頭に自分をつなぎます。

fn allocate(&mut self, data: ObjData) -> *mut Obj {
    let obj = Box::into_raw(Box::new(Obj { marked: false, next: self.head, data }));
    self.head = obj;
    self.object_count += 1;
    obj
}

Objから中身を取り出す操作自体は、生ポインタを経由するunsafeな参照外しさえ済ませてしまえば、あとは普通の安全なメソッド呼び出しです。

impl Obj {
    pub fn as_closure(&self) -> &ClosureObj {
        match &self.data {
            ObjData::Closure(closure) => closure,
            ObjData::Upvalue(_) => unreachable!("expected a closure object"),
        }
    }
    // as_upvalue / as_upvalue_mut も同様
}

こうしてunsafeの範囲を「ポインタを参照外しして&Objを得るところ」だけに絞り込み、その後のデータへのアクセスは普通の安全なRustコードとして書けるようにしています。

mark_roots——VMの状態を知っているのはVM自身

「今どのオブジェクトが生きているか」を判定するにはGCのルート(プログラムから直接たどり着ける値)を洗い出す必要があります。ルートが何であるかを知っているのはVMの内部状態そのものなので、mark-sweepの起点となるmark_rootsheap.rsではなくvm.rsに置きました(cloxでオブジェクトの確保自体はobject.c、GCサイクルの起動はvm.c/memory.cが担うのと同じ役割分担です)。

fn mark_roots(&self) {
    self.stack.iter().for_each(mark_value);
    self.globals.values().for_each(mark_value);
    self.frames.iter().for_each(|frame| mark_object(frame.closure));
    self.open_upvalues.iter().copied().for_each(mark_object);
}

値スタック・グローバル変数テーブル・呼び出し中のコールフレームが持つクロージャ・まだOpenなupvalueの一覧、この4箇所がルートです。mark_valueは値の種類を見て、ヒープを指しているものだけmark_objectに渡します。

fn mark_value(value: &Value) {
    match value {
        Value::Closure(obj) => mark_object(*obj),
        Value::Array(items) => items.iter().for_each(mark_value),
        _ => {}
    }
}

fn mark_object(obj: *mut Obj) {
    unsafe {
        if (*obj).mark() {
            return; // 既にマーク済み。循環があってもここで再帰が止まる
        }
        match &(*obj).data {
            ObjData::Closure(closure) => closure.upvalues.iter().copied().for_each(mark_object),
            ObjData::Upvalue(UpvalueData::Closed(value)) => mark_value(value),
            ObjData::Upvalue(UpvalueData::Open(_)) => {}
        }
    }
}

Obj::markは自分のマークビットを立てつつ、立てる前の状態を返します。

pub fn mark(&mut self) -> bool {
    std::mem::replace(&mut self.marked, true)
}

「既にマーク済みなら何もせず戻る」という1行が、循環参照があってもmark_objectが無限再帰しないことを保証している要です。クロージャAのupvalueがクロージャA自身を指すような循環があっても、2周目でマーク済みと判定されてそこで止まります。

sweep——アドレスを変えずに、マークされなかったものだけ解放する

マークフェーズが終わったら、リストを1周して、マークされなかったノードを解放します。マークされたノードはマークビットを次回のために下ろすだけで、アドレスは一切動かしません(コンパクション、つまり生き残ったオブジェクトを詰め直して断片化を減らす処理はしていません)。VMのあちこちが*mut Objを生ポインタのまま持ち回っている以上、GCの最中に生存オブジェクトのアドレスが変わってしまうと、それだけで全てのポインタが無効になってしまうためです。

pub unsafe fn sweep(&mut self) {
    let mut previous: *mut Obj = std::ptr::null_mut();
    let mut current = self.head;
    while !current.is_null() {
        let next = unsafe { (*current).next };
        if unsafe { (*current).marked } {
            unsafe { (*current).marked = false; } // 次回のためにマークを下ろす
            previous = current;
        } else {
            if previous.is_null() {
                self.head = next;
            } else {
                unsafe { (*previous).next = next; }
            }
            self.object_count -= 1;
            unsafe { drop(Box::from_raw(current)); }
        }
        current = next;
    }
}

Box::from_rawで確保時と対になるBoxを復元し、そのままdropに任せることでメモリを解放します。ClosureObjが持つupvalues: Vec<*mut Obj>はただのポインタの配列であり、Upvalueオブジェクトの所有権自体はこのヒープのリストだけが持っているので、クロージャを解放してもそれが指していたupvalueオブジェクトを巻き添えで壊すことはありません。

GCを起動するタイミングは、生存オブジェクト数がしきい値を超えたときです。cloxは確保済みバイト数で判定しますが、今回はオブジェクトのサイズを個別に計測していないので単純化し、オブジェクト数で判定しています。

fn maybe_collect_garbage(&mut self) {
    if self.heap.should_collect() {
        self.collect_garbage();
    }
}

fn collect_garbage(&mut self) {
    self.mark_roots();
    unsafe { self.heap.sweep(); }
    self.heap.grow_threshold(); // 次回のしきい値を「今の生存数 × 2」まで伸ばす
}

maybe_collect_garbageは、クロージャとupvalueを確保する2箇所(alloc_closurealloc_upvalue)の入り口で必ず呼ばれます。まだ影も形も無いオブジェクトを確保する直前に判定しているので、GCが走った瞬間にはこれから作ろうとしている当のオブジェクトはまだヒープに存在せず、ルートの取りこぼしを心配する必要がありません。

動作確認

examples/closures.na(パラメータ・letのローカル変数・3段ネストのupvalue捕捉を含む)を実行し、GC導入前と同じ結果になることをまず確認しました。

$ cargo run -- examples/closures.na
stack after execution: [Int(21), Int(7)]

さらに、GCの正しさを検証する専用のテストを2つ追加しています。1つ目は、しきい値を0にしてほぼ全ての確保の直前でGCを強制発火させながらプログラムを実行するテストです。

#[test]
fn closures_still_work_correctly_under_constant_gc_pressure() {
    let vm = run_with_gc_threshold(
        "fn make(x: int): int { fn inner(): int { return x; } return inner(); } \
         let a = make(1); let b = make(2); let c = make(3);",
        0,
    );
    assert_eq!(vm.stack()[0], Value::Int(1));
    assert_eq!(vm.stack()[1], Value::Int(2));
    assert_eq!(vm.stack()[2], Value::Int(3));
}

mark_rootsがルートを1つでも取りこぼしていれば、生きているはずのクロージャやupvalueがGC発火のたびに解放されてしまい、この時点でクラッシュするか誤った値を返すはずです。常時GC圧をかけた状態でも正しい結果が返ることが、マーク処理の網羅性の裏付けになります。

2つ目は、実際にメモリが回収されていることを件数で確認するテストです。

#[test]
fn gc_frees_closures_that_become_unreachable_after_return() {
    let mut vm = run_with_gc_threshold(
        "fn make(x: int): int { fn inner(): int { return x; } return inner(); } \
         let a = make(1); let b = make(2); let c = make(3);",
        1000,
    );
    let count_before = vm.heap.object_count();
    assert_eq!(count_before, 7);

    vm.collect_garbage();

    assert_eq!(vm.heap.object_count(), 3);
}

makeを3回呼ぶと、make自身のクロージャ(1個)に加えて、呼び出しごとにinnerのクロージャとupvalueが1組ずつ(合計6個)確保され、都合7個になります。しきい値を大きくして実行中に自動発火しないようにした上で明示的にcollect_garbageを呼ぶと、生存数は3個まで落ちます(この「なぜ0ではなく3なのか」は備考で説明します)。コンパイラ・VM合わせて93件の既存テストと合わせて、全てパスすることを確認済みです。

備考

なぜ3個生き残るのか——fnは常にグローバル名前解決で呼ばれる

Expr::Callは、呼び出す関数がトップレベルのものかネストしたものかを問わず、常にOP_GET_GLOBALで名前引きするようにコンパイルされます(関数とコールスタックを実装する記事で導入した、相互再帰・前方参照に対応するための設計です)。つまりfn inner(): int { ... }という宣言自体が、実行されるたびに同じグローバル名"inner"を新しいクロージャで上書きするOP_DEFINE_GLOBALになっています。makeを3回呼ぶと、globals["inner"]は3回上書きされ、最後の呼び出し(x = 3)で作られたクロージャとそのupvalueだけが、プログラム終了後もグローバル変数テーブル経由で生き続けます。これにmake自身のクロージャ(upvalue無し)を足した3個が、正しい生存数です。最初にこのテストを書いたときは「回収されるのだから0になるはず」と考えて書いたところ実際には3が返り、GCの実装ではなくテストの期待値の方が間違っていた、という経緯があります。裏を返せば、この不一致に気づけたこと自体が、GCが「本当に要らないものだけを正しく判別して回収している」ことの実証になっています。

Rust 2024 editionではunsafe fnの中身も暗黙にunsafeにならない

sweepunsafe fnとして宣言していますが、Rust 2024 edition以降はunsafe fnの本体であっても生ポインタの参照外しなどには明示的なunsafe { ... }ブロックが必要です(unsafe_op_in_unsafe_fnという2021 edition時点では警告だったルールが、2024 editionではデフォルトの挙動になりました)。C言語のcloxをそのまま読むだけでは気づけない、Rustのエディション固有の作法です。関数シグネチャのunsafe fnは「この関数を呼ぶ側が満たすべき前提条件がある」という契約を表すだけで、本体の中身が自動的に安全性チェックを免除されるわけではない、と捉えると素直に理解できます。

Rc→生ポインタへの置き換えで消えたコード

Rc<RefCell<Upvalue>>だった頃は、共有のたびにRc::clone、読み書きのたびに.borrow()/.borrow_mut()が必要でした。*mut Objは単なるポインタなのでCopyで、共有は値のコピーだけで済み、読み書きも(unsafeな参照外しさえ挟めば)直接のフィールドアクセスで完結します。実行時の借用チェックが無くなった分、安全性の保証はコンパイラの外側、つまり「ルートを取りこぼさない」「マーク前に解放しない」という設計者側の責任に移っています。

サイクルを実際に回収するテストはまだ書けない

このGCを入れた本来の動機は構造体とメソッドが作る循環参照でしたが、その構造体自体がまだ実装されていません。今回のテストで確認できているのは、通常の(循環していない)ガベージが正しく回収されることと、頻繁なGC発火の下でも既存の全機能が壊れないことの2点です。実際に循環参照を作って、それがリークせずに回収されることを示すテストは、構造体を実装した段階で追加する予定です。