✨

単純型付きラムダ計算をRustとLLVMを使ってコンパイルしてみた

2023/12/20に公開

Rust

ラムダ計算

LLVM

tech

これは大学無線部 Advent Calendar 2023の18日目の記事です。(全然無線とは関係ない内容ですが...)

https://adventar.org/calendars/9345

今回の記事は以下の記事を大いに参考にしています。OCaml(と英語)が読める人は下の記事を読む方が良いと思います。

ラムダ計算とは

ラムダ計算の項tは以下のように表されます。

t ::=     // 項
    x     // 変数
  | λx. t // ラムダ抽象
  | t t   // 関数適用(この2つのtは同じというわけではなく、項であるというのを表している)

これに加えて、今回は四則演算とプリミティブな整数も加えています。

例えば、javascriptのコード(x => y => x)(3)(2)は、ラムダ計算では(λx. λy. x) 3 2となります。

ラムダ計算はチューリング完全(チューリングマシンと同じ計算能力を持つ)であり、関数型言語の理論的基盤となっています。詳しくはWikipedia ラムダ計算を見て下さい。

単純型付きラムダ計算とは

上記のラムダ計算に(多相でない)型を付けたのが単純型付きラムダ計算です。今回は型tyは以下のように表されます。

ty ::=       // 型
    Int      // Int型
  | ty -> ty // 関数

単純型付きラムダ計算はチューリング完全ではありません。

処理の流れ

rustを入れ、リポジトリをcloneし、cargo run -- --helpとするとusageが見られます。(llvm16が必要です。)

処理の流れは次の通りです。

プログラム(文字列)を構文解析してASTにする(parser.rs)
-> α変換(後述)をする(alpha.rs)
-> 型推論をする(typeinfer.rs)
-> A正規形に変換する(anf.rs, convert)
-> closure変換をする(anf.rs, closure_conversion)
-> 巻き上げをする(anf.rs, hoisting)
-> llvm irに変換する(compile.rs)
-> llvm irを実行する(main.rs)

構文解析

(構文解析についてはあまり詳しく書きませんが)rust-pegを使っています。エラー処理を頑張るのでなければ楽に書けて使いやすいライブラリだと思います。

構文解析されると、次のように定義されるASTが生成されます。

構文解析した結果は、cargo run -- --ast "プログラム"とすると見られます。

例えば、cargo run -- --ast "(\x. \y. x)"とすると、

Abs(Variable { name: "x", id: 0 }, Abs(Variable { name: "y", id: 0 }, Var(Variable { name: "x", id: 0 })))

と表示されます。

α変換

例えば、λx. xとλy. yはどちらも引数を一つ受け取ってそれを返す関数(ラムダ抽象)です。このような、等価な束縛変数の付け替えのことをα変換(alpha conversion)と言います。束縛変数とは、λx. tのxの部分に来る変数のことです。
このα変換を施すことによって、同じ変数名であれば必ず同じ名前を表すこととなり、今後の処理が楽になります。
例えば、λx. λx. xという式であれば、λx_0. λx_1. x_1となります。

具体的な変換コードは以下を参照下さい。

環境を表すためにAlphaConvMapというリストを定義しています。ラムダ計算の場合は関数の引数は必ず一つであるため、これで環境を表すことが出来ます。一意な値としては、数字をRc<RefCell<usize>>で内部可変性パターンで持って、変数に数字を割り当てるたびにインクリメントしています。

構文解析した結果は、cargo run -- --alpha "プログラム"とすると見られます。

例えば、cargo run -- --alpha "(\x. \y. x)"とすると、

Abs(Variable { name: "x", id: 0 }, Abs(Variable { name: "y", id: 1 }, Var(Variable { name: "x", id: 0 })))

と表示されます。(xに0、yに1が割り振られています)

型推論

今回は型推論は主眼ではないので、簡潔な説明のみとなります。
Hindley-Milner型推論を実装しています。未知の変数に型変数を割り当て、等しい型同士を単一化(unify)することで型を求めます。型が正しく推論されれば、プログラムが行き詰まり状態とならず、かつ停止することが保証されます。

型は、cargo run -- --type "プログラム"とすると見られます。

実は、今回の言語をコンパイルするのにここで推論した型の情報は必要ないのですが、プログラムの正しさを検証するために型推論をしています。(型推論がなければYコンビネータが書けてチューリング完全となるのですが、プリミティブなifやBool型の実装をサボっているのであまりYコンビネータが書けても嬉しくないと思い、型推論フローを入れています。)

A正規形に変換する(A正規化)

A正規形(A-normal form)とは、次のような文法で表される形^[1]です。

Expr ::= 
    Value
  | let Variable = Value in Expr
  | let Variable(Variable*) = Value in Expr
  | let Variable = Value(Value) in Expr

Value ::=
    Variable
  | Number

要するに、式の中に現れる値が変数か数字のみになるように変換された形ということです。A正規化をすると、式が単純になるので、最適化が施しやすくなる等の利点があります。

例えば、(λx. λy. x + y) 2 3であれば、次のように変換されます。

let f_2(x_0) = 
  let f_3(y_1) = 
    let z_4 = x_0 + y_1 in
    z_4 in
  f_3 in
let y_5 = f_2(2) in
let y_6 = y_5(3) in
y_6

A正規形の定義は以下のようにしました。

Value::Globalはグローバルな関数のポインタを表します。
ANFのそれぞれの要素の1つ目の引数は、束縛される変数名を表しています。このあたりは、ANFのDisplayの実装をみると分かりやすいと思います。

ANFsは、anfsがletで定義される部分のVectorで、valueは結果の式、levelはlet式の深さです。

A正規化のコードは次のようになります。

ラムダ抽象の場合のみ新しいANFsを作り、そのANFsにラムダ抽象内部の式の変換結果を書き出しています。(L214)
その他の場合は渡されたANFsにそのまま書き込んでいます。

A正規形は、cargo run -- --anf "プログラム"とすると見られます。

closure変換

次にするのはclosure変換です。現状のA正規形では、関数の内部に自由変数が現れてしまっています。自由変数とは、束縛変数でない変数のことです。例えば(λx. λy. x + y) 2 3のA正規形

let f_2(x_0) = 
  let f_3(y_1) = 
    let z_4 = x_0 + y_1 in
    z_4 in
  f_3 in
let y_5 = f_2(2) in
let y_6 = y_5(3) in
y_6

では、f_3関数の内部に現れるx_0がf_3から見て自由変数となっています。この先llvm irに変換するためには、関数がすべてグローバルに定義されている必要があるのですが、f_3に自由変数があるためf_3をそのまま外に出すことが出来ません。

そこで、closure変換(closure conversion)という処理を行い、関数内部から自由変数を消します。
closureとは、日本語では関数閉包と呼ばれるもので、関数と環境を組にしたものです。このclosureに含まれる環境として自由変数の値を持たせることで、関数内部の自由変数を無くすことが出来ます。

具体例を見てみましょう。先の(λx. λy. x + y) 2 3のA正規形

let f_2(x_0) = 
  let f_3(y_1) = 
    let z_4 = x_0 + y_1 in
    z_4 in
  f_3 in
let y_5 = f_2(2) in
let y_6 = y_5(3) in
y_6

をclosure変換すると、

let f_8(env_7, x_0) = 
  let f_10(env_9, y_1) = 
    let x_0 = env_9[1] in
    let z_4 = x_0 + y_1 in
    z_4 in
  let f_3 = (@f_10, x_0) in
  f_3 in
let f_2 = (@f_8) in
let f_11 = f_2[0] in
let y_5 = f_11(f_2, 2) in
let y_12 = y_5[0] in
let y_6 = y_12(y_5, 3) in
y_6

が得られます。((@f_8)は@f_8を表すのではなく、@f_8一つのみからなる組を表していることに注意して下さい。)
ここで、f_10に注目すると、変換前では自由変数であったx_0は、f_10では渡された環境から取り出されています。そして、f_8の返り値が、f_10とx_0の組(closure)となっています。また、関数呼び出しの形態が少し変わっています。closureから0番目の値(=関数)を取り出し、その関数にclosure自身と元の引数を与えて呼び出しています。

closure変換のコードは以下のようになります。

変換されるのは、関数定義と関数適用の場合です。
まず関数定義の場合は、let var(args) = funbody_anfsとして、funbody_anfsに現れる自由変数を計算します。その後、closure変換後のfunbody_anfsの先頭に、環境から計算された自由変数の値を取り出し、束縛するコードを追加します。そして、argsの先頭に新しく生成した環境の変数を追加し、関数を定義します。その後に、計算された自由変数の先頭に定義した関数を追加し、組を作ります。

関数適用の場合は、closureから関数を取り出し、その関数にclosure自身と引数を渡します。

closure変換の結果は、cargo run -- --closure "プログラム"とすると見られます。

巻き上げ

関数内部にある関数定義を外側に移動します。トップレベルにある関数定義以外の定義はmain関数の中に入れます。

変換された結果はcargo run -- --hoist "プログラム"とすると見られます。

例えば、cargo run -- --hoist "(\x. \y. x + y) 2 3"を実行すると、

let f_10(env_9, y_1) =
  let x_0 = env_9[1] in
  let z_4 = x_0 + y_1 in
  z_4

let f_8(env_7, x_0) =
  let f_3 = (@f_10, x_0) in
  f_3

let main() =
  let f_2 = (@f_8) in
  let f_11 = f_2[0] in
  let y_5 = f_11(f_2, 2) in
  let y_12 = y_5[0] in
  let y_6 = y_12(y_5, 3) in
  y_6

となります。

見ての通り、無駄の多い変換となっています。closure変換の最適化については、今後調べていきたいですね。

LLVM IRに変換

LLVMとは、LLVM IRという型付き高級アセンブリのようなものを最適化し、各マシン向けの機械語を出力してくれるとても便利なものです。ここでは、巻き上げ後のA正規形をLLVM IRに変換します。

このLLVM IRに変換するために、inkwellというライブラリを使っています。これはllvm-sysというローレベルなLLVMのC APIバインディングをより安全にしたものです。

今回は関数の引数と返り値はすべてi64型(64bitの整数型)で表します。これは多相型の実装でよく使われる方法で、もともと64bitで表される型はそのまま、そうでない型は参照を取得してポインタ型をi64型とします。

コード全体は次のようになります。

inkwellの使い方については省略して、それぞれの定義がどのようにLLVM IRに変換されるのかついて書いていきます。

関数定義
let fun_name(args) = fun_bodyとします。関数定義はトップレベルにのみ現れます。関数の型は、全て(i64, i64) -> i64で表されます。(もともと1引数関数しか無く、かつすべての関数がclosure変換されて環境を受け取るようになっているため。)
変数名とLLVM IRの対応を持っておくためのHashMapを作り、各引数名とIRの対応を追加します。その後、fun_bodyをLLVM IRに変換します。fun_bodyの中に含まれる定義を変換した後に、fun_bodyのvalueを変換し、それを返り値とします。

例えばλx. x + 1という関数(ラムダ抽象)であれば、次のように変換されます。

define i64 @f_4(i64 %0, i64 %1) {
entry:
  %z_2 = add i64 %1, 1
  ret i64 %z_2
}

関数適用
let var = fun_var(args)とします。まず、環境からfun_varに対応する値を取り出します。その値の型をi64から関数のポインタ型に変換します。この変換を行う命令がinttoptrというもので、ポインタ型が64bitであればCPUでは何も行われない命令となります。
そして、argsをLLVM IRに変換します。その後、関数ポインタに、変換されたargsを渡します。その結果をvarに束縛し、環境に追加します。

例えば、let y_3 = f_6(f_1, 2)であれば、次のように変換されます。

%ptr4 = inttoptr i64 %f_6 to ptr
%y_3 = call i64 %ptr4(i64 %ptr1, i64 2)

四則演算
let var = val1 op val2とします。まずval1、val2を変換し、opの種類に応じて演算をします。その結果をvarに束縛し、環境に追加します。

例えば、let z_4 = x_0 + y_1であれば、次のように変換されます。

%z_4 = add i64 %x_0, %1

組
let var = (tuple)とします。まず、tupleのそれぞれの要素をLLVM IRに変換します。その後、mallocを呼び出して組用のメモリを確保し、ポインタを取得します(tuple_ptr)。確保するメモリの大きさは8*(tupleの長さ)byteです。
その後は、変換されたtupleのそれぞれの要素をtuple_ptrに割り当てていきます。まずgetelementptr命令を使って割り当て先のポインタを取得します。得られたポインタにtupleの変換された要素を割り当てます。割り当てるにはstore命令を使います。
すべてのtupleの要素を割り当てた後、tuple_ptrをptrtoint命令を用いてi64型に変換します。ptrtoint命令はinttoptr命令とは逆の命令となります。変換されてi64型となった値をvarに束縛し、環境に追加します。

例えば、let f_3 = (@f_8, x_0)であれば、次のように変換されます。

%f_3 = call ptr @malloc(i64 16)
%ptr = getelementptr i64, ptr %f_3, i32 0
store i64 ptrtoint (ptr @f_8 to i64), ptr %ptr, align 4
%ptr1 = getelementptr i64, ptr %f_3, i32 1
store i64 %1, ptr %ptr1, align 4
%ptr2 = ptrtoint ptr %f_3 to i64

射影
let var = tuple_var[index]とします。まず、環境からtuple_varに対応する値を取り出します。その値の型をi64からi64のポインタ型にinttoptr命令を用いて変換します。そのポインタからindex番目の要素のポインタをgetelementptr命令を用いて取得し、load命令を用いて値を取り出します。その値をvarに束縛し、環境に追加します。

例えば、let x_0 = env_7[1]であれば、次のように変換されます。

%ptr = inttoptr i64 %0 to ptr
%ptr1 = getelementptr i64, ptr %ptr, i32 1
%x_0 = load i64, ptr %ptr1, align 4

以上のルールに従って、まずmain以外の関数を変換し、その後main関数を変換します。main関数の型は() -> i64となっています。

LLVM IRはcargo run -- --llvm "プログラム"とすると見られます。

LLVM IRを実行

JITコンパイルをして、main関数を実行しています。

まとめ

今回は単純型付きラムダ計算をLLVMを使ってコンパイルしました。closure変換や関数型言語のコンパイラの中間表現について知りたかったため、実装することが出来て良かったです。

脚注

あくまでこの記事の中では、です。 ↩︎

GitHubで編集を提案