🚀

自作言語でLLVM JITによるマクロが動くまで

2021/12/26に公開

LLVM

compiler

tech

この記事は言語実装 Advent Calendar 2021の12日目の記事です。12日はとっくに過ぎているのですが、カレンダーが空いていたので書かせて頂きました。

自作言語llrlには、LLVM JITを用いたマクロ機能が備わっています。この機能の実装過程での課題を振り返ります。

以下の記事もご参照ください。

マクロの背景: llrlの概要

マクロに関連する機能を中心的に、llrl言語自身について紹介します。
llrlはLisp-likeな構文の自作プログラミング言語です。特に特定の用途に特化しているわけでもない汎用(general-purpose)プログラミング言語ですが、Lispのように、プログラムがS式で記述されるのが特徴的です。いくつか例を見てみます。

(let* ([a 3]
       [b 5]
       [c (+ a b)])
  (println! "a + b = " c))

; 関数定義は (function signature body ...) で行う
(function (fibonacci-number x)
  (if (< x 3)
    1
    (+ (fibonacci-number (- x 1))
       (fibonacci-number (- x 2)))))

プログラムがS式で記述される一方、基本的な言語機能はStandardML系の言語に近いものとなっています。プログラムはHindley-Milner型推論による静的型付けがなされます。プログラム中に型シグネチャを明示的に記述することもできます。

; 以下のような関数定義は...
(function (square x)
  (* x x))

; 以下のような型に推論・決定される:
(function (square x) {(forall A) (-> A A) (where (Mul A))}
  (* x x))
; ここで、
; (forall ..) は型パラメタ
; (-> A A)    はこの関数の型 (ここではAを引数に取ってAを返す)
; (where ..)  は制約 (ここではAはかけあわせられる性質を持つ (Mul A) ことを要求)

llrlは代数的データ型の定義とパターンマッチを備えています。例えば builtin モジュールには、OCamlやRustでよく知られているデータ型 Option, Result が定義されています。

; 値データ型の定義は (value-data type-constructor value-constructor ...) で行う
(value-data (Option A) ; (Option A) 型の値は
  none                 ; none か
  (some A))            ; (some A) のいずれか

(value-data (Result T E)
  (err E)
  (ok T))

Option, Result に関連して、モジュール std/option や std/result から基本的な操作が関数として提供されています。パターンマッチの用例に丁度良いので抜粋すると、例えば map 関数は以下のように定義されています。

(function (option/map f x) {(forall A B) (-> (-> A B) (Option A) (Option B))}
  (match x           ; パターンマッチ: 入力 x が
    [(some (let x))  ;   (some (let x)) にマッチするなら
      (some (f x))]  ;     関数適用してsomeで包んで返す
    [none            ;   none にマッチするなら
      none]))        ;     noneを返す

(function (result/map f x) {(forall TA TB E) (-> (-> TA TB) (Result TA E) (Result TB E))}
  (match x
    [(ok (let x))
      (ok (f x))]
    [(err (let e))
      (err e)]))

llrlのマクロの概要

llrlのマクロは、Lisp系言語の伝統的なマクロに近いものとなっています。マクロは 入力としてS式を取り、マクロ展開後のS式またはエラーを返します 。一例として、引数の式を完全に無視して、常に式 #t に展開されるマクロは以下のように定義できます。

(macro (always-true s)
  (ok '#t))

その実、マクロの定義は関数定義とほとんど変わりません。マクロの本体式に制限は無く、通常の関数の本体式と同じようにプログラムを記述できます。関数定義と異なるのは、関数として見たマクロの型が常に (-> (Syntax Sexp) (Result (Syntax Sexp) String)) 固定であることです。この型は「入力にS式 (Syntax Sexp) を取り、返値として展開結果 (Syntax Sexp) またはエラー String を返す」と読むことができ、これは上のマクロの説明とそのまま対応しています。

Result 型については上に書いた通り、Rust等の同名の型と同等な働きをするデータ型でした。 Result 型と同じく、 Syntax および Sexp 型についても builtin モジュールに定義されています。

(builtin-type (Syntax A) "syntax"
  [(syntax A) "syntax"])

(value-data Sexp
  (sexp:integer Bool U64)
  (sexp:fpnumber F64)
  (sexp:bool Bool)
  (sexp:symbol String)
  (sexp:string String)
  (sexp:char Char)
  (sexp:cons (Syntax Sexp) (Syntax Sexp))
  sexp:nil
  (sexp:use CapturedUse))

Syntax はメタ情報(ソースコード中の位置など)を内包する組み込みデータ型で、ここでは (syntax <data>) がS式のツリー構造の中に常に挟まるとだけ考えてもらえれば良いです。 Sexp 型はまさしく、S式を構成する構造が単に代数的データ型のコンストラクタとして定義されています。いくつか対応例を見てみましょう。

プログラム	対応する `(Syntax Sexp)` の値
`123`	`(syntax (sexp:integer #f 123))`
`#t`	`(syntax (sexp:bool #t))`
`(#t)`	`(syntax (sexp:cons (syntax (sexp:bool #t)) (syntax sexp:nil)))`
`(#t #f)`	`(syntax (sexp:cons (syntax (sexp:bool #t)) (syntax (sexp:cons ...))))`

(syntax (sexp:...)) を (s:...) として省略するとよりわかりやすくなります:

プログラム	対応する `(Syntax Sexp)` の値
`123`	`(s:integer #f 123)`
`#t`	`(s:bool #t)`
`(#t)`	`(s:cons (s:bool #t) (s:nil))`
`(#t #f)`	`(s:cons (s:bool #t) (s:cons (s:bool #f) (s:nil)))`

このようなプログラムとデータの対応から、 プログラマは、この (Syntax Sexp) 型のデータを分解・変換して返すプログラムを記述することでマクロを実装できる といえます。 (Syntax Sexp) の操作自体は他のデータ構造の操作と変わりません。

なお、実践的には、 s/match や準クォート (quasiquote) などを用いることで、明示的な (Syntax Sexp) の分解・構築処理を書かずに… value-data Sexp のデータ構造を意識せずにマクロを実装することができます。

; lambdaマクロの定義例: s/matchによってS式の分解を、準クォート(quasiquote)によってS式の構築を行っている
(macro (lambda s)
  (s/match s
    [(_ ,args ,@body)
      (ok
        (let ([tmp-f (gensym)])
          `(let ([(,tmp-f ,@args) ,@body])
            ,tmp-f)))]
    [_
      (err "Expected (lambda (arg ...) body ...)")]))

; 例えば、
; (lambda (x y) (+ x y)) が
; (let ([(<tmp-f-generated-sym> x y) (+ x y)]) <tmp-f-generated-sym>) と展開される
; ※ <tmp-f-generated-sym> はgensymによって生成されたユニークなシンボル

この s/match や準クォート (quasiquote) は、言語レベルでは標準ライブラリに定義されている通常のマクロです (2-s-match.llrl , 3-quasiquote.llrl)。

そのため、処理系のレベルでやることは特に変わらず、処理系は(1)マクロ定義を必要に応じて実行可能な状態にし、(2)マクロがプログラム中で使用されたときに、S式を入力 (Syntax Sexp) として与えて実行し、結果をS式として再解釈できればマクロ機能を実現できるといえます。

前置きが長くなりましたが、マクロ機能の実現のために何が必要かが見えてきました。

(1) 使用したいマクロを呼び出せるように準備する

前述のとおり、llrlのマクロは、通常の関数と同じように本体式に任意のプログラムを記述できます。マクロは通常の関数と同様に意味解析され、通常の関数と同様にLLVMのFunctionへとコンパイルされます。言い換えると、 マクロは通常の関数と同様に、意味解析を経てコンパイルが完了するまで利用可能となりません。 そのため、プログラムを解析してコンパイルしていく順番が重要になります。

(function hello
  (always-true 123)) ; マクロの使用: このとき always-true は利用可能か?

(macro (always-true s) ; マクロの定義
  (ok '#t))

これについては、様々な方法が考えられます。

プログラムを常に上から順に解釈していく?

素朴なやり方の一つに「プログラムを常に上から順に解釈していく」という方法が考えられます。処理系は上から順にプログラムを解釈していき、未知の定義を使用しようとした場合は即座にエラーとする方法です。この方法は処理系が考えることは非常に少なくなりますし、トップレベルの定義をランタイムに実行される手続きのシーケンスと見たときに「上から下へ」実行されるのと一致していて理に適っています。一方、定義の使用の解決を静的に行う言語では、これはかなり強い制約です。

(function (even? x)
  (if (= x 0)
    #t
    (odd? (- x 1)))) ; odd? が定義される前に使用されている

(function (odd? x)
  (if (= x 0)
    #f
    (even? (- x 1))))

「上から順に解釈」ではこのようなプログラムは許容されません。筆者はトップダウンにコードを書くのが好きで、このような相互参照や前方参照を許容したかったため ^[1]、「プログラムを常に上から順に解釈していく」という方法は取りませんでした。

依存関係に基づいた順に解釈していく?

llrlの処理系は、前方参照や相互参照を許容するために、意味解析を何段階かのパスに分けて行います。

モジュール内の定義の表を作り
モジュール内の全ての使用をその表から解決し
その上で型推論などの解析を行う

特に3つ目のパスについては、各定義の順番に関わらず一貫性のある結果が得られるように、内部的に定義同士の依存関係を計算して順に型推論を行って...といった処理を行っています (型システムもこのような型チェックが可能なものを採用しています)。この方法...「依存関係に基づいた順に解釈していく」をマクロについても適用できないかと考えられます。

; (2) hello が依存する always-true を処理したので hello を処理する
(function hello
  (always-true 123))

; (1) このモジュールの他の定義に依存しない always-true を処理する
(macro (always-true s)
  (ok '#t))

この方法はうまく適用できそうに思われますが、うまく機能しない場合があります。以下の例を考えます:

(value-data Vec
  (vec: F64 F64 F64))

(getter vec: vec/x vec/y vec/z)

ここで、 getter はデータの各要素を取り出す関数を定義するマクロと考えてください (実装例)。このマクロは展開の結果 新たな関数定義が得られます が、依存関係の解析のためには 既に定義と使用の解決がなされている 必要があります。上に挙げたようなパスを順番に1度ずつ適用していく方法では意味解析が完結しなくなってしまいます。

マクロの使用をうまく解析できるようにマクロを使用できる場所を制限したり、依存関係やその後の意味解析を部分的に遅延したり、といった解決の糸口は勿論考えられますが、いずれにしても処理系の実装が複雑になりそうでした。

マクロは別のモジュールで定義されたもののみ利用できる

これまでの選択肢のトレードオフを踏まえ、llrlでは、言語仕様としてより強い(しかし、シンプルな)制約を設けて、これに基づいて「使用したいマクロを呼び出せるように準備しておく」という要求を満たせるようにしました。その制約が 「マクロは別のモジュールで定義されたもののみ利用できる」 というものになります。

llrlにおいて、ソースコードは各ファイルがそれぞれモジュールに対応します (1ファイル = 1モジュール)。モジュールの仕様はECMAScriptのそれに近く、定義を export したり、別のモジュールから定義を import したりできるというものです。モジュールの集合もモジュール内の定義の集合と同様に、 import 宣言からモジュール間の依存関係が解析され、依存される側のモジュールから順番に処理されます。

llrlにおいてモジュールの意味解析は他のモジュールと完全に独立しています。モジュール a に依存するモジュール b があったとき、まずモジュール a の意味解析が 完全に行われてから 、モジュール b の意味解析が行われます。これは、前述のような、モジュール内の定義の集合がいっせいに何段階かのパスを通して少しずつ解析されるのとは対照的です。

このようなモジュールの仕様から、「マクロは別のモジュールで定義されたもののみ利用できる」という制限を設けることで、 プログラム中で使用されているマクロが準備できているかを意味解析のレベルで考慮する必要がなくなります 。そのマクロが同一モジュールにある場合は仕様上利用が許可されておらず、別モジュールから import したものであれば、その別モジュールは既に意味解析が完了しているためです。

うまく説明できているか怪しいのですが、 言語仕様上の不便さの裏には時にこのようなトレードオフがある という話でした... プログラミング言語の利用者には普段見えない形で、この制約をもっと緩めることは可能でしょうが、llrlでは言語仕様の簡単さや実装しやすさを考えてこのような制約を設けることにしました。

厳密な話をすると、モジュール間の依存を発生させる import もまたプログラム中に記述されるため、意味解析の対象のはずです。

(import "builtin" Sexp U64 Number.U64)
(import "std/prelude/stage-0" _)

これについてllrlは、各モジュールに行う主要な意味解析とは別に、事前に import だけ集め、まずモジュール間の依存関係を解析するというステップを挟んでいます。当然これはマクロ展開よりも前に行われるので、llrlには マクロで import 宣言を生成することができない という制限があります。

(2) マクロがプログラム中で使用されたときにマクロを実行する

「(1) 使用したいマクロを呼び出せるように準備する」の話と比べると、(2)ではひたすら泥臭い実装が続きます。

(1)でマクロからコード生成するまでの準備(意味解析)は完了するようにしたので、そのままLLVMのFunctionを生成して、LLVMのJIT実行エンジンでそのFunctionを実行することでマクロが動作します。具体的には、マクロ呼び出しを発見するたびに、そのマクロおよびマクロに関連する実装のLLVM-IR生成がまだなされていない場合は新たにLLVMのモジュールを生成してLLVMの実行エンジンに加え、最後にマクロに対応するLLVM Functionを実行します。

二つの世界のデータの相互変換

一連の手続きは、LLVMのモジュールを細切れに作り続ける以外はLLVMのKaleidoscope Tutorialで解説されているそのままですが、マクロ実行においては マクロ側(つまりllrl側)のS式の内部表現 と ホスト実装側(Rust)のS式の内部表現 を相互変換する必要があります。

// Rust側のS式の表現
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Hash, new)]
pub struct Sexp {
    pub loc: SourceLocation,
    pub rep: SexpRep,
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Hash)]
pub enum SexpRep {
    Integer(bool, u64),
    FPNumber(OrderedFloat<f64>),
    Symbol(Cow<'static, str>),
    String(Cow<'static, str>),
    Char(char),
    List(Vec<Sexp>),
    ListLike(Box<(Sexp, Vec<Sexp>, Sexp)>),
    Bool(bool),
    Native(Native),
}

; llrl側の表現
(derive (Eq Ord DebugDisplay Hash) data (S A)
  (s: SourceLocation (SRep A)))

(derive (Eq Ord DebugDisplay Hash) value-data (SRep A)
  (s-rep:integer Bool U64)
  (s-rep:fp-number F64)
  (s-rep:bool Bool)
  (s-rep:symbol String)
  (s-rep:string String)
  (s-rep:char Char)
  (s-rep:cons (S A) (S A))
  s-rep:nil
  (s-rep:pure A))

これについて、効率の良い方法を考えたりはしたものの (内部表現を一致させるとか遅延させるとか)、2つの世界のギャップが大きく (特にllrl側はBoehm GCによるアロケーションを前提としている)、諦めてマクロ呼び出しのたびに愚直に相互変換する形としました。

Rust側に、 llrlのコンパイル後の内部表現と一致するように注意深くレイアウトされた データ型を用意します。以下は Sexp 型に対応するデータ型の例です。コンパイル後の内部表現と一致させるため、Rustで使用するのは珍しいであろう union などが出てきます。

#[derive(Clone, Copy)]
#[repr(C)]
pub struct RtSexp {
    tag: u64,
    body: RtSexpBody,
}

#[derive(Clone, Copy)]
#[repr(C)]
union RtSexpBody {
    integer: RtSexpInteger,   // 0
    fpnumber: RtSexpFPNumber, // 1
    bool: RtSexpBool,         // 2
    symbol: RtSexpSymbol,     // 3
    string: RtSexpString,     // 4
    char: RtSexpChar,         // 5
    cons: RtSexpCons,         // 6
    nil: RtSexpNil,           // 7
    use_: RtSexpUse,          // 8
}

このようなデータとRust側の本来のデータとの相互変換を都度行うことになります。

特にボックス化の起点となる RtSyntax (llrlの Syntax 型に対応します) では、相互変換時には以下のような処理が伴います。

llrl側への変換時はBoehm GCの GC_malloc を用いてアロケート、 ptr::write
Rust側への変換時はllrlで管理されている領域 (= Boehm GCの管理領域) から ptr::read

llrlのセルフホスティング実装においては、内部表現が一致しているのでこのような相互変換を要しないというのが少し面白いところでしょう。

終わりに

マクロ機能の実装は、次々と課題が浮かび上がってきて大変でした。production-readyなプログラミング言語にあるマクロ機能にも色々と制約があったりしますが、その理由の一端が掴めたような気がします。次はScala 3あたりのマクロに触れてみたいと思います。

脚注

OCamlの let rec .. and .. のような構文をS式の中で見い出せなかったのもある ↩︎