Schemeコンパイラで、末尾再帰のクロージャをループに変更する

2014-01-15

考えていた、Schemeで末尾再帰のコードをループに変更する方法を実装した。考えとしてはシンプルだけど、実装するには大きな変更が必要で、174コミット、15日も費やしてようやく実装できた(まだバグがあるかもしれないけど…)。やったことをメモしておきたいと思う。

1.モチベーション

Schemeの言語仕様にはループのための構文はなく、その代わりに末尾呼び出し最適化を保証している。末尾呼び出し最適化があればスタックを大量に消費せずに繰り返しを実現できる。

実際にSchemeでループをさせるには、例えばnamed letを使って、

(let loop ((i 10)
(acc 0))
(if (> i 0)
(loop (- i 1) ; <= 末尾呼び出し
(+ acc i))
(print acc)))

のようになる。このコードはマクロによって展開されて

((lambda (loop)
(set! loop (lambda (i acc)
(if (> i 0)
(loop (- i 1)
(+ acc i))
(print acc))))
(loop 10 0))
nil)

のようになる。

しかしこのコードを字面通りナイーブにコンパイルすると、ループ用のクロージャの生成と、変数への破壊的代入が発生してしまう。これらはどちらも実行時にメモリ確保が必要で、パフォーマンス的によろしくないので、できれば避けたい。

そこで文脈を解析して、これらを解決する。

2.構文木を作成するコンパイラへの変更

クロージャをループに変換してもいいかどうかを調べるために、ループ用の関数の変数(上の例ではloop)がどのように使われているか知る必要がある。3impで示されるような、Schemeのプログラムから直接バイトコードを出力する形式のコンパイラだと変数の使われ方を調べるのが難しいので、一旦S式からスコープ情報を持った構文木に変換して、その構文木を辿ってコンパイルする方式に変更した。

スコープはlambda式が出てくるたびに現在のスコープを親とする新しいスコープを作成する。スコープにそれぞれの変数がどのように使われているかを構文木の構築時に格納していく。

構文木のノードは、例えば(lambda (loop) body...)とかを単に(:LAMBDA scope (loop) body-node body-code)のように付加情報を追加したものに変換していく。

3.クロージャをループに変換できる条件

前回クロージャをループに変換できるための条件を考えたが、実際にはもっと色々条件が必要だった。 lambda式のset!は、変数が導入されたスコープの地の文で1回だけ 関数としてだけ使われていて、他の関数への引数や、戻り値として使われてはいない lambda式の内部では末尾再帰のみ(末尾じゃない位置での再帰がない) lambda式の内部の直下からだけ呼び出されていて、さらに深い位置からではない ループ開始のための、地の文での呼び出しが1回だけ(これは末尾じゃなくてもよい)

4.コンパイル時のスコープ変更

構築された構文木のコンパイルでは、ループとして使われるlambda式のset!は完全に出力を省略する。そしてループ開始のための呼び出し部分に、lambda式の本体を出力する。ここでスコープの変更が必要になる。

元のソースで説明すると、

((lambda (loop)  ; <= スコープA
(set! loop (lambda (i acc) ; <= スコープB
(if (> i 0)
(loop (- i 1)
(+ acc i))
(print acc))))
(loop 10 0)) ; <= C
nil)

ループ本体用クロージャのスコープBは、構文木を作成した時点ではスコープAを親とする別スコープなんだけど、ループとして出力する場合、このset!自体の出力がまるまる省略される。そしてCの位置でのループ開始でその本体をコンパイルして出力するのだけど、スコープBは出力されなくなってしまうので、新たにスコープAをパラメータ(ここでは(i acc))で拡張したスコープで本体をコンパイルする必要がある。

このスコープの変更は結構トリッキーで、スコープを変えて構文木を再構築するわけにも行かず、スコープBのメンバへの破壊的代入で対応している。さらにスコープBのlambda式の本体ノードに対してもスコープBの変更を適用する必要がある。ここは破壊的代入だらけで汚いし、バグもいっぱい出た。本当にうまく動いているのかよくわかってない。

5.ループ継続のジャンプ

ループを継続するための末尾再帰の関数呼び出しは、ジャンプに変更できる。それには末尾呼び出し最適化と同様に、スタック操作が必要となる。新しいVMのオペコード(LOOP)を追加して対応した。

6.結果

今までのコンパイラでの出力:

(GREF nil
PUSH
EXPND 1
BOX 0 ; ループ用関数を格納する変数をボックス化
FRAME (CONST 0
PUSH
CONST 10
PUSH
LREF 0
PUSH
;; ループ用関数のクロージャの生成
CLOSE 2 1 (FRAME (CONST 0
PUSH
LREF 0
PUSH
GREF >
APPLY 2)
TEST (FRAME (LREF 0
PUSH
LREF 1
PUSH
GREF +
APPLY 2)
PUSH
FRAME (CONST 1
PUSH
LREF 0
PUSH
GREF -
APPLY 2)
PUSH
FREF 0
UNBOX
TAPPLY 2) ; ループ用関数を末尾再帰で呼び出し
LREF 1
PUSH
GREF print
TAPPLY 1)
LSET 0
APPLY 2)
SHRNK 1
HALT)

ボックス化とクロージャ生成が行われていた。

ループに変換する版の出力:

(GREF nil
PUSH
EXPND 1
;; ボックス化が省略されている
CONST 0
PUSH
CONST 10
PUSH
EXPND 2 .
;; クロージャの生成が省略されている
#0=(FRAME (CONST 0
PUSH
LREF 0
PUSH
GREF >
APPLY 2)
TEST (FRAME (LREF 0
PUSH
LREF 1
PUSH
GREF +
APPLY 2)
PUSH
FRAME (CONST 1
PUSH
LREF 0
PUSH
GREF -
APPLY 2)
PUSH
LOOP 2 1 . #0#) ; 末尾再帰が直接ジャンプに変更されている
FRAME (LREF 1
PUSH
GREF print
APPLY 1)
SHRNK 2
SHRNK 1
HALT))

ボックス化とクロージャ生成が省略され、また末尾再帰が直接のジャンプに変更されている。これにより、単純なループがメモリ確保をせずに実行できるようになった。

7.その他

  • ループに変更した場合、生成されるコードには自己参照が発生するので、コンパイルした結果を出力したり、デバッグで表示させようとする場合、write/ssのような、無限ループに陥らない対策が必要
  • またそのコンパイルして出力されたコードを読み込むにはリーダーが自己参照の読み込みに対応する必要がある
  • 定義している地の文とかlambda式内部の地の文という条件がlet式を挟んだだけで成り立たなくならないように、3imp 4.7.2 Direct Function Invocations の実装が必要

単なるループを実装するためだけにこんだけ苦労するのって一体…素直にループ構文導入すればいいじゃんと思うけど、これもラムダ計算の自由度を手に入れるため…と我慢する。