ocaml の多相比較関数の最適化の問題について調べた

もう一月前の話ですが…忘れる前に書いておきます。

ocaml の多相関数が、より効率のよい単相な関数に置き換えられるのは、全適用の形でかかれている場合のみ、というちょっとした罠があります。詳細はこちら→ http://d.hatena.ne.jp/camlspotter/20100130/1264852627
十分に型付けされているにも関わらず最適化されないのはおかしい、という主張は確かに最もですね。というわけでなぜそんな理不尽なことになっているのか調べてみました。

bytecomp/translcore.ml の transl_prim 関数は、名前の通りプリミティブを置き換えます。この関数は、プリミティブとそれに対して適用されている引数（ノード）を受け取って、適切な関数（ノード）を返します。
もう面倒くさいのでコード貼りますが、

let transl_prim prim args =
  try
    let (gencomp, intcomp, floatcomp, stringcomp,
         nativeintcomp, int32comp, int64comp,
         simplify_constant_constructor) =
      Hashtbl.find comparisons_table prim.prim_name in (* テーブルから引いてくる *)
    begin match args with
    (* 略。ヴァリアントみたいな int でないけど intcomp を返すものの処理 *)
    | [arg1; arg2] when has_base_type arg1 Predef.path_int
                     || has_base_type arg1 Predef.path_char ->
        intcomp
    | [arg1; arg2] when has_base_type arg1 Predef.path_float ->
        floatcomp
    | [arg1; arg2] when has_base_type arg1 Predef.path_string ->
        stringcomp
    | [arg1; arg2] when has_base_type arg1 Predef.path_nativeint ->
        nativeintcomp
    | [arg1; arg2] when has_base_type arg1 Predef.path_int32 ->
        int32comp
    | [arg1; arg2] when has_base_type arg1 Predef.path_int64 ->
        int64comp
    | _ ->
        gencomp
    end
  with Not_found ->
    (* 以下、比較関数以外のプリミティブの処理 *)

とまあこんな具合なので、部分適用の場合には多相のままになるのは自明ですね。部分適用の場合でも単相関数返していい気がするんですが、どうなんでしょうね。ここは試しに実装してみなかったので、よく分かりません。
分かる人はここみただけでオチが分かると思います。ボクは分からなかったので、もう少し調べました。

勿論関数適用でない、つまり識別子でしかない状態の比較関数も置き換える必要があります。これにはもう全てを諦めた transl_primitive という別の関数があって、

let transl_primitive p =
  let prim =
    try
      let (gencomp, _, _, _, _, _, _, _) =
        Hashtbl.find comparisons_table p.prim_name in
      gencomp
    with Not_found ->
      (* 以下、比較関数以外のプリミティブの処理 *)

こうなっている。全くやる気がないですね。

型情報ちゃんと使えばなんとかなるでしょ、ということで下のようなコードを書いてみました。

let transl_primitive p type_expr =
  let is_int_compate_type lhs rhs =
    let equal_int_type = equal_type Predef.path_int in
    equal_int_type lhs && equal_int_type rhs in
  let prim =
    try
      let (gencomp, intcomp, _, _, _, _, _, _) =
        Hashtbl.find comparisons_table p.prim_name in
      match type_expr with (* option なのは、この関数使ってる場所で型情報取るの大変な場所があったから手抜き… *)
      | Some expr ->
        begin match expr.desc with
        | Tarrow (_, lhs, { desc = Tarrow (_, rhs, ret, _) }, _) when is_int_compate_type lhs rhs -> (* int -> int -> ??? なら *)
          intcomp
        | _ -> gencomp
        end
      | _ -> gencomp
    with Not_found ->
      (* 以下、比較関数以外のプリミティブの処理 *)

ウルトラ手抜きですが、こいつで

let c : int -> int -> int = compare

というようなコードを -dparsetree する。するとどうなる？

i:~/src/ocaml % ocaml -dparsetree c.ml
(* 略 *)
Pexp_constraint
expression (c.ml[1,0+28]..c.ml[1,0+35])
  Pexp_ident "compare"
(* 略 *)

ダメでしたー o(*^^*)o ワーイ

もうちょっと調べてから気付いたんですが compare を（正確には compare という名前の識別子のノードを）処理してるころには、まだ unification は行われてないので、compare にまで型の情報が伝搬されてないんですね。時間なくて unification のタイミングまでは調べられなかったんですが、何にせよそういうことらしい。誰かどのタイミングでやってるのか教えてください。typedtree の変換終えてからになるのかな？

ということで unification 終わったあとでツリー舐め直せば普通に最適化することは可能だと思います。何でやってないのかはよく分かりません。誰か最適化オプションとして実装しませんか？おしまい。