Template Haskellでコンパイル時にプログラムを走らせる。

Template Haskellを利用してコンパイル時に定数計算を行う方法とIO処理をする方法について書きます。

コンパイル時定数計算といえば、C++でのテンプレートメタプログラミングを利用した方法などが有名ですが、同じ事をしてみます。

簡単な例として、フィボナッチ数列を使います。
まず、Fib.hsというファイルを用意します。定義は普通の定義にします。

module Fib where

fib :: Int -> Int
fib n | n < 2     = 1
      | otherwise = fib(n-1) + fib(n-2)

fib 35を計算してみます。

% cat foo.fib
import Fib
main = print $ fib 35

実行結果。

% ghc --make foo.hs                                                                                                                  
% time ./foo                                                                                                                         
14930352
./foo  2.40s user 0.06s system 99% cpu 2.465 total

時間を図っているのはコンパイル時に計算出来ている事を簡単に確認する為です。
ただ、愚直な実装でfib 35が2.4秒ってのは意外と早い気がしますね。さすがにCにくらべると10倍近く遅いです。

Template Haskellでコンパイル時に計算を行うには以下のように書きます。

{-# OPTIONS_GHC -fth #-}
import Language.Haskell.TH
import Fib

main = print $( let x = fib 35 
                 in [| x |] )

最初に$( ... )の部分がコンパイル・実行されるので、fib 35の結果が計算されたあと,Quasi-quote bracketsによって数値リテラルを表す内部表現に変換されます。
そして、それがprintの引数部分に展開されて埋め込まれるという流れです。
当然コンパイル時間は長くなります。

% ghc --make foo.hs                                                                                                                  [2 of 2] Compiling Main             ( foo.hs, foo.o )
Loading package base ... linking ... done.
Loading package template-haskell ... linking ... done.    <- ここでfib 35が実行されている
Linking foo ...
% time ./foo                                                                                                                         
14930352
./foo  0.00s user 0.00s system 62% cpu 0.006 total

$( ... )の中のコンパイルが先ですので、$( ... )内で参照する関数、変数は外部からインポートしているか、$( ... )内部で定義したものでなければならないというところが多少面倒ですが、Fib.hsのコードがそのままコンパイル時にも利用できるので簡単です。

こんな感じでrunIOを使うと

{-# OPTIONS_GHC -fth #-}
import Language.Haskell.TH
import Fib

main = print $( let x = fib 35 
                 in do runIO $ putStrLn "Calculating fibonacci ..."
                       [| x |] )

コンパイル時にメッセージを出せます。

% ghc --make foo.hs                                                                                                                 
[1 of 2] Compiling Fib              ( Fib.hs, Fib.o )
[2 of 2] Compiling Main             ( foo.hs, foo.o )
Loading package base ... linking ... done.
Loading package template-haskell ... linking ... done.
Calculating fibonacci ...   <- オリジナルログメッセージ
Linking foo ...

これは楽しい。