Mokkosuで型推論器を作ってみた

はじめに

MokkosuやF#、OCaml、Haskellなどの言語には型推論器が実装されており、ユーザが型を一切明示しなくても関数の引数や戻り値の型をコンパイラが推論してくれます。ここではMokkosuを使って簡単な型推論器を作ります。

式の構文

まず、型推論する式の構文を定めます。今回対象とする言語は、ラムダ計算 + 整数 + 足し算の言語としたいと思います。Mokkosuではそのような言語の抽象構文木は代数的データ型を使って簡単に表すことができます。

type Expr = Int(Int)
          | Add(Expr, Expr)
          | Var(String)
          | Fun(String, Expr)
          | App(Expr, Expr)
;

例えば、Mokkosuの\f -> \x -> f x + 1は、

Fun ("f", Fun ("x", Add (App (Var "f", Var "x"), Int 1)))

のようになります。

型の構文

次に、型の構文を式と同様に定めます。この言語には、整数型と関数型があります。型推論を行うためにはこの2つに加えて型に関する変数(型変数)を追加する必要があります。

type Type = TInt
          | TFun(Type, Type)
          | TVar(Int, Ref< Maybe<Type> >)
;

今回は型変数は型変数を一意に識別するための整数と、実際の型を格納するリファレンスセルの組とします。リファレンスセルはMaybe<Type>型を要素として持ち、型変数に何も代入されていないときはNothing型tが代入されているときはJust tで表すことにします。

型推論

型推論を行う関数は以下のようになります。

fun typeinf tenv = {
  ~Int(_) -> TInt;
  ~Add(e1, e2) ->
    let t1 = typeinf tenv e1 in
    do unify t1 TInt in
    let t2 = typeinf tenv e2 in
    do unify t2 TInt in
    TInt;
  ~Var(x) -> from_just (lookup x tenv);
  ~Fun(x,e) ->
    let t1 = new_tvar () in
    let tenv2 = (x, t1) :: tenv in
    let t2 = typeinf tenv2 e in
    TFun(t1, t2);
  ~App(e1,e2) ->
    let t1 = typeinf tenv e1 in
    let t2 = typeinf tenv e2 in
    let t3 = new_tvar () in
    do unify t1 (TFun (t2, t3)) in
    t3;
};

型推論の関数は型環境(変数名と型の連想リスト)と式を受け取ってその式の型を返します。

まず、式がInt(_)の形の時は、整数型なのでTIntを返します。

式がAdd(e1, e2)の形の時は、以下のように処理します。

    # e1 の型を推論する
    let t1 = typeinf tenv e1 in
    # e1 の型がTIntであることを確認する
    do unify t1 TInt in
    # e2 の型を推論する
    let t2 = typeinf tenv e2 in
    # e2 の型がTIntであることを確認する
    do unify t2 TInt in
    # Add(e1,e2)全体の型はTInt
    TInt;

ここでunify関数は以下の型の関数で型t1と型t2を等しくする働きがあります。簡単に処理を説明すると、t1とt2を比較して型が等しければなにもせず、等しくない場合はエラーにします。片方が型変数の時は、変数に型の代入を行います。unify関数は次の節で定義するのでそこで詳しく解析してください。

式が変数Var(x)の時は、型環境から対応する型を探索して返します。

式がFun(x,e)の時は、以下のように処理します。

    # 新しい型変数t1を作る
    let t1 = new_tvar () in
    # 変数xの型をt1として式eを型推論する
    let tenv2 = (x, t1) :: tenv in
    let t2 = typeinf tenv2 e in
    # Fun(x,e)の型は(t1 -> t2)型
    TFun(t1, t2);

式がApp(e1,e2)の時は、以下のように処理します。

    # e1の型を推論する
    let t1 = typeinf tenv e1 in
    # e2の型を推論する
    let t2 = typeinf tenv e2 in
    # 新しい型変数t3を用意する
    let t3 = new_tvar () in
    # t1 = (t2 -> t3) となるように代入を行う
    do unify t1 (TFun (t2, t3)) in
    # App(e1,e2)の型はt3
    t3;

単一化 (unification)

単一化とは2つの型を等しくする処理のことです。単一化は以下のように定義できます。説明はコメントに書いてあります。

fun unify t1 t2 =
  match (t1, t2) {
    # t1 = TInt, t2 = TInt の場合もともと等しいので何もしない
    (~TInt, ~TInt) -> ();
    # 両方が関数型の場合、引数と戻り値それぞれを等しくする
    (~TFun(t11, t21), ~TFun(t12,t22)) ->
      do
        unify t11 t12;
        unify t21 t22;
      end;
    # 両辺が同じ型変数の場合、何もしない。
    (~TVar(n1,r1), ~TVar(n2,r2)) ? n1 == n2 -> ();
    # 片方が型変数でなおかつ具体的な型が代入されている場合、
    # その中身ともう片方の型を等しくする。
    (~TVar(n1,r1), _) ? is_just !r1 -> unify (from_just !r1) t2;
    (_, ~TVar(n2,r2)) ? is_just !r2 -> unify t1 (from_just !r2);
    # 片方が型変数でなおかつ何も代入されていない場合、
    # その変数にもう片方の型を代入する。
    # ただしここでoccur n1 t2でないことを確認する(後述)。
    (~TVar(n1,r1), _) ->
      if occur n1 t2 -> error "occurs error"
      else r1 := Just t2;
    (_, ~TVar(n2,r2)) ->
      if occur n2 t1 -> error "occurs error"
      else r2 := Just t1;
    # それ以外の場合(たとえばt1が整数型でt2が関数型など)は型エラー
　  _ ->
      error "unify error"
  };

occur n t関数は型tの内部にIDnの型変数が含まれていないことを確認する関数です。例えば、

(\x -> x x)

のような式は、xの引数にx自信が現れるため、有限長の型で表現できません。そのためにエラーとしてはじいています。これを出現検査とよびます。

出現検査 (occurs check)

型tの中に型変数nが現れないことを確認する処理です。以下のように定義します。

fun occur n = {
  # TIntは型変数ではない
  ~TInt -> false;
  # 関数型のときはどちらかに型変数nが含まれていればtrueを返す
  ~TFun(t1, t2) -> occur n t1 || occur n t2;
  # 型変数の場合
  ~TVar(m, r) ->
    # n が含まれていればtrue
    if n == m -> true
    else
      # そうでない場合は中身に対して検査する。
      match !r {
        ~Nothing -> false;
        ~Just(t) -> occur n t
      }
};

プログラム全体

これまで出てきた関数を組み合わせてプログラムを作ります。

# Mokkosu単相型推論

type Expr = Int(Int)
          | Add(Expr, Expr)
          | Var(String)
          | Fun(String, Expr)
          | App(Expr, Expr)
;

type Type = TInt
          | TFun(Type, Type)
          | TVar(Int, Ref< Maybe<Type> >)
;

let var_count = ref 0;

fun new_tvar () =
  do incr var_count in
  TVar(!var_count, ref Nothing);

fun occur n = {
  ~TInt -> false;
  ~TFun(t1, t2) -> occur n t1 || occur n t2;
  ~TVar(m, r) ->
    if n == m -> true
    else
      match !r {
        ~Nothing -> false;
        ~Just(t) -> occur n t
      }
};

fun unify t1 t2 =
  match (t1, t2) {
    (~TInt, ~TInt) -> ();
    (~TFun(t11, t21), ~TFun(t12,t22)) ->
      do
        unify t11 t12;
        unify t21 t22;
      end;
    (~TVar(n1,r1), ~TVar(n2,r2)) ? n1 == n2 -> ();
    (~TVar(n1,r1), _) ? is_just !r1 -> unify (from_just !r1) t2;
    (_, ~TVar(n2,r2)) ? is_just !r2 -> unify t1 (from_just !r2);
    (~TVar(n1,r1), _) ->
      if occur n1 t2 -> error "occurs error"
      else r1 := Just t2;
    (_, ~TVar(n2,r2)) ->
      if occur n2 t1 -> error "occurs error"
      else r2 := Just t1;
    _ ->
      error "unify error"
  };

fun typeinf tenv = {
  ~Int(_) -> TInt;
  ~Add(e1, e2) ->
    let t1 = typeinf tenv e1 in
    do unify t1 TInt in
    let t2 = typeinf tenv e2 in
    do unify t2 TInt in
    TInt;
  ~Var(x) -> from_just (lookup x tenv);
  ~Fun(x,e) ->
    let t1 = new_tvar () in
    let tenv2 = (x, t1) :: tenv in
    let t2 = typeinf tenv2 e in
    TFun(t1, t2);
  ~App(e1,e2) ->
    let t1 = typeinf tenv e1 in
    let t2 = typeinf tenv e2 in
    let t3 = new_tvar () in
    do unify t1 (TFun (t2, t3)) in
    t3;
};

fun type_to_string = {
  ~TInt -> "int";
  ~TFun(t1, t2) -> "(" ^ type_to_string t1 ^ " -> " ^
                   type_to_string t2 ^ ")";
  ~TVar(n, r) ? is_just !r -> type_to_string (from_just !r);
  ~TVar(n, r) -> to_string n;
};

let test e = type_to_string (typeinf [] e);

let test01_ref = \x -> x + 3;
let test01 = Fun ("x", Add (Var "x", Int 3));

let test02_ref = \f -> \x -> f (x + 0);
let test02 = Fun ("f", Fun("x", App (Var "f", Add (Var "x", Int 0))));

let test03_ref = \f -> \g -> (f g) 0;
let test03 = Fun ("f", Fun("g", App (App (Var "f", Var "g"), Int 0)));

do println <| test test01;
do println <| test test02;
do println <| test test03;

実行結果

上のプログラムを実行すると、

(int -> int)
((int -> 4) -> (int -> 4))
((6 -> (int -> 8)) -> (6 -> 8))

と表示されます。4や6、8は型が代入されていない変数です。Mokkosuでいうところの[?]型、OCamlの'_a型に相当します。とりあえず正しい形の型が推論されていることが分かると思います。

拡張

今回は最低限の型推論器を実装しました。このプログラムはいろいろと拡張が可能です。例えば以下のようなものが挙げられます。

• 他の基本型(Bool, String, Floatなど)を入れる。
• タプルやEither型を入れる。
• 多相型に対応させる。

ぜひ挑戦してみてください。