OCamlで128行で作るJVMバイトコードコンパイラ

一昨年に出版されたImplementing Programing Languages <http://www.digitalgrammars.com/ipl-book/> という本（以下IPL）を読んでいました。

この本は名前の通りインタプリタ/コンパイラの実装についての本ですが、とてもいい本です。何がまずいいかというと、薄い（そして安い）。207ページしかありません。コンパイラ作りたいとなったらドラゴンだのタイガーだのと格闘しなければならないという先入観がありますのでこの薄さは画期的です。

内容の深さという意味では他の本に譲る点はあるのでしょうが、一通りのことが実践できるようになるまでを要領よく解説していて、かつ最初は読み飛ばしてもいい理論部分にもわざわざ印をつけてくれているので、おそらくコンパイラ作成志望者にとっての最短コースを示してくれる本だと思います。

飛ばし読みの成果としてJVMバイトコードを生成する簡単な言語処理系を作ってみることにしました。

文法の定義

文法の定義はIPLでも使っているBNFCという文法定義言語/ツールを使います。ツールはUbuntuではapt-getで入れられます。

言語名は仮にTauとしました。データは整数のみで、条件分岐と四則演算とローカル変数定義と関数呼び出しを持ち、しかしループは持ちません。関数定義はトップレベルに1つのみ、引数も1つのみに限定しました。

Tau.cfというファイル名で文法を定義します。

Func. Func ::= "fun" Ident "(" Pat ")" "=" Exp ;

Add. Exp ::= Exp "+" Exp1 ;
Sub. Exp ::= Exp "-" Exp1 ;
Mul. Exp1 ::= Exp1 "*" Exp2 ;
Div. Exp1 ::= Exp1 "/" Exp2 ;
Atm. Exp2 ::= Atom ;

coercions Exp 2 ;

Let. Atom ::= "let" Pat "=" Exp "in" Exp ;
Cnd. Atom ::= "if" Exp "then" Exp "else" Exp ;
Int. Atom ::= Integer ;
Var. Atom ::= Ident ;
Exp. Atom ::= "(" Exp ")" ;
App. Atom ::= Ident Exp ;

VarPat. Pat ::= Ident ;

これをBNFCにかけると、この文法から大体想像がつくような抽象構文木データ型とパーサージェネレータ用のソースが生成されます。簡単！
この文法だとFunc, Exp, Atom, Patという直和データ型ができて、Add, Sub...などのコンストラクタができるわけです。（※追記: この文法定義は関数適用の優先順位が一般的に期待されているのと異なる点やshift-reduce conflictが起こる点でちょっといまいちだとあとで分かったのですが、とりあえずそのままにしておきます。またcoercionを定義していれば Exp. Atom ::= "(" Exp ")" ; のルールは不要でした。）

パーサージェネレータは昔はコンパイラコンパイラと呼ばれていましたので、BNFCはコンパイラコンパイラコンパイラということになるのだそうです。

IPLではHaskelとJavaが説明に使われていますが、BNFCツール自体はOCaml(ocamlyacc/ocamllex)やその他の言語にも対応していますのでここではOCamlを使います。

bnfc -ocaml -m Tau.cf

-mはMakefileを生成するオプションです。

構文木からJavaアセンブラへの変換

構文木に対する操作を行うコードは自分で書く必要がありますが、一緒にSkel[文法名].ml（ここではSkelTau.ml）というスケルトンのソースも生成されるのでそれを元に作ればよいです。

IPLを参考にしつつ、構文木をJVMバイトコードに変換するためのソースを書きました。

module Compile = struct

open AbsTau

type label = string
type instr =
    Label of label | Ldc of int | Iadd | Isub | Imul | Idiv | Ifeq of label
  | Goto of label | Ireturn | Return | Aload of int | Invokenonvirtual of string
  | Iload of int | Istore of int | Invokestatic of string

type env = int ref * (string, int) Hashtbl.t
let extend_env name (top, tbl) = (Hashtbl.add tbl name (!top); incr top)
let lookup_env name (_, tbl) = Hashtbl.find tbl name
let empty_env () = (ref 0, Hashtbl.create 10)

let label_index = ref 0
let new_label () = (label_index := (!label_index + 1); "LABEL" ^ string_of_int (!label_index))

let emit = function
  Label l -> print_endline @@ l ^ ":"
| Ldc i -> print_endline @@ "\tldc " ^ string_of_int i
| Iadd -> print_endline "\tiadd"
| Isub -> print_endline "\tisub"
| Imul -> print_endline "\timul"
| Idiv -> print_endline "\tidiv"
| Ifeq l -> print_endline @@ "\tifeq " ^ l
| Goto l -> print_endline @@ "\tgoto " ^ l
| Ireturn -> print_endline "\tireturn";
| Return -> print_endline "\treturn";
| Aload i when i <= 3 -> print_endline @@ "\taload_" ^ string_of_int i
| Aload i             -> print_endline @@ "\taload " ^ string_of_int i
| Invokenonvirtual m -> print_endline @@ "\tinvokenonvirtual " ^ m
| Iload i when i <= 3 -> print_endline @@ "\tiload_" ^ string_of_int i
| Iload i             -> print_endline @@ "\tiload " ^ string_of_int i
| Istore i when i <= 3 -> print_endline @@ "\tistore_" ^ string_of_int i
| Istore i             -> print_endline @@ "\tistore " ^ string_of_int i
| Invokestatic m -> print_endline @@ "\tinvokestatic Tau/" ^ m ^ "(I)I"

let rec compileFunc (x : func) env = match x with
    Func (Ident name, VarPat (Ident param), exp) -> (
      extend_env param env;
      print_endline @@ ".method public static " ^ name ^ "(I)I";
      print_endline ".limit locals 256";
      print_endline ".limit stack 256";
      compileExp exp env;
      emit Ireturn;
      print_endline ".end method")
and compileExp (x : exp) env = match x with
    Add (exp0, exp) -> (compileExp exp0 env; compileExp exp env; emit Iadd)
  | Sub (exp0, exp) -> (compileExp exp0 env; compileExp exp env; emit Isub)
  | Mul (exp0, exp) -> (compileExp exp0 env; compileExp exp env; emit Imul)
  | Div (exp0, exp) -> (compileExp exp0 env; compileExp exp env; emit Idiv)
  | Atm atom -> compileAtom atom env
and compileAtom (x : atom) env = match x with
    Let (VarPat (Ident name), exp0, exp) -> (
      extend_env name env;
      compileExp exp0 env;
      emit @@ Istore (lookup_env name env);
      compileExp exp env)
  | Cnd (exp0, exp1, exp) ->
    let false_label = new_label () in
    let true_label = new_label () in (
      compileExp exp0 env;
      emit (Ifeq false_label);
      compileExp exp1 env;
      emit (Goto true_label);
      emit (Label false_label);
      compileExp exp env;
      emit (Label true_label))
  | Int n -> emit (Ldc n)
  | Var (Ident name) -> emit @@ Iload (lookup_env name env)
  | Exp exp -> compileExp exp env
  | App (Ident name, exp) -> (
    compileExp exp env;
    emit @@ Invokestatic name)

let compile func = (
  print_endline @@ ".class public Tau";
  print_endline ".super java/lang/Object";
  print_newline ();
  print_endline ".method public ()V";
  emit (Aload 0);
  emit (Invokenonvirtual "java/lang/Object/()V");
  emit Return;
  print_endline ".end method";
  print_newline ();
  compileFunc func (empty_env ()))
end

open Lexing

let parse (c : in_channel) : AbsTau.func =
    ParTau.pFunc LexTau.token (Lexing.from_channel c)

let main () =
    let filename = Sys.argv.(1) in
    let channel = open_in filename in
    try
        Compile.compile (parse channel)
    with BNFC_Util.Parse_error (start_pos, end_pos) ->
        Printf.printf "Parse error at %d.%d-%d.%d\n"
            start_pos.pos_lnum (start_pos.pos_cnum - start_pos.pos_bol)
            end_pos.pos_lnum (end_pos.pos_cnum - end_pos.pos_bol);
;;

main ()

ちょっとスタイルが命令的すぎる感じもしますが、IPLでもコード生成のところは凝ったことしないで命令的に書いたほうがいいよと言っているんですね。

JVMバイトコード出力と書きましたが、実際にはJasminというアセンブラ用のアセンブリを出力しています。 関数はJVM世界の静的メソッドに対応します。JVMのメソッドは何らかのクラスに属する必要がありますのでTauというクラスに属させることにしました。

このソースをMakefileをいじってビルドします。

サンプルプログラム

おなじみの階乗計算はこの言語ではこのように書きます。

fun f(x) =
  if x then
    x * f (x - 1)
  else 1

これを先ほど作成したコンパイラにかけるとJasminソースができます。

$ ./tau fact.tau
.class public Tau
.super java/lang/Object

.method public ()V
        aload_0
        invokenonvirtual java/lang/Object/()V
        return
.end method

.method public static f(I)I
.limit locals 256
.limit stack 256
        iload_0
        ifeq LABEL1
        iload_0
        iload_0
        ldc 1
        isub
        invokestatic Tau/f(I)I
        imul
        goto LABEL2
LABEL1:
        ldc 1
LABEL2:
        ireturn
.end method

一度ファイルに落としてからjasmin（これもapt-getで入れられます）でアセンブルするとクラスファイルが出来上がります。

$ ./tau fact.tau > Tau.j
$ jasmin Tau.j
$ javap -cp . Tau
Compiled from "Tau.j"
public class Tau {
  public Tau();
  public static int f(int);
}

ドライバ

このクラスはpublis static int f(int)という静的メソッドを含むだけのクラスなので計算結果を表示するためにはドライバが必要です。これは面倒なのでJavaで書きましょう。

public class Driver {
  public static void main(String[] args) {
    System.out.println(Tau.f(Integer.parseInt(args[0])));
  }
}

これをコンパイル、実行すると階乗関数を呼び出してくれます。

$ javac -cp . Driver.java
$ java -cp . Driver 10
3628800

一応電卓で検算すると…あってました。

ローカル変数も使ってみます。

$ cat a.tau
fun f(x) = let a = x + 3 * 4 in 5 + 6 * a
$ ./tau a.tau > Tau.j
$ jasmin Tau.j
$ java -cp . Driver 2
89

できてますね。

BNFC生成コードを除外して、自分で書いた部分文法定義とコード生成箇所だけを合計すると126行しか書いていませんでした。なかなかいい感じだと思います。

$ cat Tau.cf Compile.ml | wc -l
126