RustのarenaでGleamの3年来の性能問題を解決する
(giacomocavalieri.me)- Gleamフォーマッタの再帰的な
Document構造で繰り返し発生していた個別のヒープ割り当てを、arenaベースの参照に置き換え、3年間残っていた性能改善の課題を解決 Box<Self>をarenaに格納された&Selfへ置き換え、Rustのライフタイム検査により、arenaが破棄された後に内部データが参照されないことを保証- 言語キーワードやカンマのように繰り返し現れる
Documentを数百個一度だけ割り当てて再利用することで、プリティプリンタの実行時間が13msから9.8msに短縮され、24%高速化 - ソースの読み込みやパースを含む
gleam format全体の実行時間も13%短縮され、最大メモリ使用量は8.4MBから7.6MBへ約10%減少 - arenaを渡すように関数と呼び出し側を広範囲に修正する必要があり、
+2963/-1032規模の手作業が必要だったが、繰り返される個別割り当てを減らす方法で速度とメモリ使用量の両方を改善
再帰的なDocument構造で発生していた割り当てコスト
- Gleamのプリティプリンタは、コードをどのように出力し、行幅制限を超えたときにどこで分割するかを表す再帰的な**
Documentデータ構造**を使っている - 主なvariantは、出力する文字列を保持する
String、改行の前後の文字列を定義するBreak、複数の文書をまとめるGroup、改行時にインデントを増やすNestなどで構成される - リスト
[1, 2]もString、Break、Nest、Groupを組み合わせて表現する- 現在の行に収まるなら
[1, 2]のように1行で出力する - 行幅制限を超える場合は、指定された
Breakに従って複数行に分割してインデントし、末尾のカンマも出力する
- 現在の行に収まるなら
Nestのように別のDocumentを内部に持つvariantは、従来Box<Self>を使っていたため、入れ子になった文書ごとにヒープへ個別に割り当てる必要があり、この処理が実行時間のかなりの部分を占めることがあった- Gleamプリティプリンタの基礎アルゴリズムはStrictly Prettyの論文で確認できる
Boxをarena参照に置き換える
- 従来の
Nest(Box<Self>)をNest(&'doc Self)に変更し、arenaに保存された別のDocumentを参照する構造へと改めた - 参照を保持するため、文字列ライフタイム
'stringとは別に、文書データのライフタイム'docをDocumentに追加した - 実装には
typed_arenaクレートを使用allocで値をarenaに保存すると、その値への参照が返される- 割り当てられたデータはarenaが生存している間利用できる
- arenaがスコープを外れると、内部データもまとめて破棄される
alloc(&self, value: T) -> &mut Tが返す参照はarenaより長く生きられず、Rustの借用チェッカーがarena破棄後の不正な参照を防ぐ
繰り返し使う文書を一度だけ割り当てて再利用
- arenaは他の文書を参照するためだけでなく、コード全体で繰り返し現れる
Documentをキャッシュする用途にも使われている String("fn")、String("pub")、String("type")のような言語キーワードを表す文書を毎回新しく作る必要がなくなった- リスト項目の間で使う
Break { unbroken: ", ", broken: "," }のような文書も、一度割り当てた後に再利用する - 従来は延々と
Boxに入れる必要があった数百個の小さな文書を、それぞれ一度だけ割り当てればよくなった
単純なAPIと広範囲な呼び出し側の変更
- arena API自体は単純だが、既存コードの大部分を移行する過程では反復的な修正が必要だった
- 以前は
Box::newを呼ぶだけだった関数も、今ではデータを割り当てるためのarena引数を追加で受け取る必要がある format_listのような関数は、リスト項目をフォーマットする下位関数にも同じarenaを渡し、入れ子の文書をBox::newではなくarena.allocで生成するように変わった- 数日かけて慎重に検索・置換を進め、最終的に
+2963/-1032規模のpull requestとして作業を完了した - 反復的で手作業中心の作業であり、この過程でLLMは使われなかった
速度と最大メモリ使用量の改善
- 実際のGleamプロジェクトである
squirrelをフォーマットした際、プリティプリンタの実行時間は13msから9.8msへ短縮された- プリティプリンタ単体で見れば24%の性能向上にあたる
gleam formatはフォーマット以外にも、プロジェクトのソースコードの読み込みやパースを行うため、プリティプリンタは全体処理の一部にすぎないが、全体の実行時間も13%短縮された- 最大メモリ使用量は8.4MBから7.6MBに減少し、約10%の削減となった
- arenaによって個別のヒープ割り当て回数を減らした結果、フォーマッタの実行速度が向上し、メモリ使用量もあわせて減少した
1件のコメント
Lobste.rsの意見
ピークメモリ使用量をさらに10%減らしたいなら、参照の代わりに不透明ハンドルを返す別のアリーナを検討する価値がある。ハンドルを64ビットではなく16ビットや32ビットにできるため、小さな割り当てが多い場合に効果が大きい
https://donsz.nl/blog/arenas/ によると、これをサポートするクレートは
compact_arenaだけだが、要件に合わなければ自分で実装するのもそれほど難しくないOption<T>のリストを、要素ごとに1ビット弱の存在フラグと、実際に存在する要素のTに分けて保存するその代わり、参照型として
&Option<T>ではなくOption<&T>を返す。より一般的には、構造的再帰の中に参照を置くのではなく、参照をたどって構造的再帰を行う形になるアリーナはもっと広く使われる価値があり、ハンドルやスマートポインタの代わりに参照を返すと、Rust では再帰的列挙型のパターンマッチが非常にやりやすくなる。通常は
if let MyAst::Add(MyAst::Multiply(a, b), c) = expr {}のように書こうとしても、途中の box やポインタのせいでできないが、参照なら nightly 機能なしで透過的にパターンマッチできる。すべての参照が1つのアリーナを指していれば、木全体をつなぐライフタイムも一貫して保てるこの方式は、別の Wadler方式 の pretty-print ライブラリである
prettyがサポートしているアプローチとかなり似ている。ただし、複数のメモリ割り当て方式をサポートしているため、API はやや重く感じられるかもしれない記事で引用されていた 論文 の著者が、Brunswick のごく普通の IT 企業で働いているとは驚きだ