OCamlの末尾再帰について

この記事はMLアドベントカレンダー22日目の記事です。

プロローグ

autotakerはHaskellのList.sortが遅いので嘆いていた。
あまりに遅いので簡単なコードを書いて実験することにした。

gist082402e29104b2f21a6d

$ wc test300000.in
  300000  300000 2994897 test300000.in
$ head test300000.in
-145530267
165962464
79995549
-41622317
-133297703
38688159
191031379
-188084614
-187672271
-148432431
./ListSort < test300000.in > test300000.out
[2015-12-22 06:48:09.644106 UTC] begin
[2015-12-22 06:48:09.71299 UTC] parse done
[2015-12-22 06:48:10.610641 UTC] sort done
[2015-12-22 06:48:10.610854 UTC] Sorting: 0.897651s
[2015-12-22 06:48:10.707304 UTC] output done
[2015-12-22 06:48:10.707538 UTC] Elapsed Time: 1.063198s

30万個の整数列をソートするのに900msもかかるのだ。

比較のためにC++でもソートしてみた。

gist1b7a9a6f6b9c0f446ae5

$ g++-5 -std=c++11 -O3 sort.cpp -o sort
$ time ./sort < test300000.in > test300000.out
Sorting: 22ms

real	0m0.151s
user	0m0.133s
sys	0m0.010s

C++のstd::sortは爆速でわずか22msでソートが完了した。なんとHaskellのList.sortの40倍以上の速さである。これはいくらなんでも速すぎる。
調べたところstd::sortはIntro sortと呼ばれるQuick sortを少し賢くしたようなアルゴリズムで実装されている一方でList.sortはmerge sortに毛が生えたようなアルゴリズムで実装されているらしい。アルゴリズムが違うならまあ仕方ない。

ところで、標準ライブラリのソートにマージソートが採用されていてよくHaskellと比較される関数型言語があるではないか。
そうOCamlである。OCamlコンパイラはGHCに比べてそれほどアグレッシブな最適化をするわけでもないし、さすがにOCamlよりは速いだろうと予想してコードを書いてみた。

OCamlのソートのバグ？

さて、ここまでが長い前振りで、ここから今回の本題に入ろう。
まず、次のようなOCamlのコードを書いて実験してみた。

gistebed443ea61424295eca

$ ocamlopt sort.ml -o sort-ml
$ time ./sort-ml < test300000.in > test300000.out
Sorting: 0.14416s

real	0m0.874s
user	0m0.393s
sys	0m0.470s

全体の実行時間が遅いのは主に入出力のせいなので気にしないこととして、
OCamlのList.sortは144msであり、Haskellの6倍以上速い。
これは意外な結果だったが、念のため出力があっているかを確認したときに興味深いことがわかった。

$ diff test300000.ans test300000.out | head
2,3d1
< -214745540
< -214744494
14d11
< -214729150
18d14
< -214725140
21d16
< -214723392
37d31

なんと正しくソートできていないではないか！
この短く単純なコードのどこに私はバグを仕込んでしまったのだろう。

まずは入力を正しく読めているか確認した。

let main _ = 
    let ps = read_ints [] in
    let start_t = Sys.time() in
    Printf.fprintf stderr "length: %d\n" (List.length ps);
    let qs = List.sort compare ps in
    let end_t = Sys.time() in
    Printf.fprintf stderr "Sorting: %.5fs\n" (end_t -. start_t);
    output_ints qs;;

./sort-ml < test300000.in > test300000.out
length: 261056
Sorting: 0.14787s

どうやらバグはread_intsにあるらしい。

let rec read_ints acc =
    try
        let v = read_int() in
        read_ints (v::acc)
    with
        _ -> List.rev acc

このコードを書いた時の私の気持ちはこうだ。

1. 入力はデカイので末尾再帰にする必要がある。アキュムレータ引数を用意しよう。
2. read_intのドキュメントにはファイルの末尾に達したら、End_of_file例外が投げられるとかいてある。tryで包んで例外を捉えよう。
3. 例外の名前打ち込むの面倒だし、使い捨てのコードだから_でマッチさせてもいいよね！

3つ目が私の怠慢であることは認めるが、上の二つは正しそうに思えないだろうか。

私は自身の怠慢を認め、期待する例外の名前を正しく書くことにした。

let rec read_ints acc =
    try
        let v = read_int() in
        read_ints (v::acc)
    with
        End_of_file -> List.rev acc

$ ocamlopt sort.ml -o sort-ml
$ ./sort-ml < test300000.in > test300000.out
Fatal error: exception Stack_overflow

!!!!!Fatal error: exception Stack_overflow!!!!!
読者の中にはもう気づいた人もいるだろうが、実はこのread_intsは末尾再帰になっていないのだ。
なぜかというと、read_intsにはread_ints(v::acc)を呼び出したあとに例外処理のコードがあるので、例外をキャッチした時にどのハンドラに渡すかを保存するためにスタックが必要になるためだ。

正しく末尾再帰にするためには例外をoption型に変換してやらないといけない。

let rec read_ints acc =
    let r = try Some (read_int()) with 
        End_of_file -> None in
    match r with
        Some(v) -> read_ints (v::acc)
        | None -> List.rev acc

このコードは多少不恰好だが、定数サイズのスタックしか消費しない。

gist764a6d75741e2f807bc5

$ ocamlopt sort1.ml -o sort1-ml
$ time ./sort1-ml < test300000.in > test300000.out
length: 300000
Sorting: 0.18463s

real	0m0.972s
user	0m0.436s
sys	0m0.530s
$ diff test300000.out test300000.ans

こうしてバグは修正された。めでたしめでたし。

後日談というか、今回のオチ

こうしてOCamlコードのバグは直ったわけだが、実験の結果、OCamlのList.sortは184msであり、やはりHaskellのList.sortがクソ遅いのはアルゴリズムのせいではなく、実装が悪いのだと考えられる。
競技Haskellerとしてこんな状況は許されないのでメモリ効率が悪いListの代わりにVectorを使ってマージソートを実装した。
gistdb303c4aae8e1ce9776b

$ ghc -O2 VecSort.hs
[1 of 1] Compiling Main             ( VecSort.hs, VecSort.o )
Linking VecSort ...
$ ./VecSort < test300000.in > test300000.out
[2015-12-22 07:58:53.725504 UTC] begin
[2015-12-22 07:58:53.805494 UTC] parse done
[2015-12-22 07:58:54.00343 UTC] sort done
[2015-12-22 07:58:54.003704 UTC] Sorting: 0.197936s
[2015-12-22 07:58:54.092547 UTC] output done
[2015-12-22 07:58:54.092758 UTC] Elapsed Time: 0.367043s

このように198msと、OCamlと大差ない時間になった。