RULESによるコンパイル時プログラミング

これはHaskell Advent Calendar 2013の(3+π)日目の記事です。

(3 + pi)や(quot 7 8)のような単純な定数式は、ghc -Oが行なう定数畳み込みによってコンパイル時に計算される。uncurry (*) (3, max 5 2)のようなやや複雑な式も、インライン展開してから定数畳み込みをすることでやはりコンパイル時に整数リテラルにまで簡約される。

これは一見万能だが、再帰的な関数が一つでもあると何もできなくなる。GHCが再帰関数をインライン化しないからだ。(sum [1])ですら実行時のループにコンパイルされる*1。

どうしてもコンパイル時に計算してほしい関数がある場合はどうしたら良いか。一つの方法はTemplate Haskell(ja)を使うことだが、特別な構文を使わなければいけないこと、-fwarn-unused-binds(ja)をはじめとして警告がいくつか効かなくなることなど副作用がある。この記事は、もう一つの方法、書き換え規則(ja)を使ってこの種のコンパイル時計算を実装する術を紹介する。

例として、二つの整数の最大公約数を計算する関数を考える。Haskellで普通に書くと以下のようになる。

-- | 二つの整数の最大公約数を求める
-- Preludeのgcdと区別するためにgcd'という名前にしておく
gcd' :: Int -> Int -> Int
gcd' x 0 = x
gcd' x y = pgcd (abs x) (abs y)
{-# INLINE gcd' #-}

-- | 非負の整数と正の整数の最大公約数を求める
pgcd :: Int -> Int -> Int
pgcd x y
  | q == 0 = y
  | otherwise = pgcd y q
  where
    q = mod x y

関数gcd'は非再帰的なので、INLINEプラグマ(ja)を付けておけば間違いなくインライン展開される。問題は再帰的に定義されたpgcdだ。これをRULESプラグマで何とかすることを考えよう。すぐに思い付くのは、以下のような規則を使ってコンパイル時に強引に再帰を起こすことだ*2。

{-# RULES
"pgcd" forall x y. pgcd x y =
    let q = rem x y
    in if q == 0
        then y
        else pgcd y q
  #-}

これは要するに、pgcdの定義をそのままRULEとして書き直したものだ。残念なことにこれはうまく行かない。試してみると、GHCがパニックを起こす。

ghc: panic! (the 'impossible' happened)
  (GHC version 7.6.3 for x86_64-unknown-linux):
        Simplifier ticks exhausted
    When trying RuleFired Class op ==
    To increase the limit, use -fsimpl-tick-factor=N (default 100)
    If you need to do this, let GHC HQ know, and what factor you needed
    To see detailed counts use -ddump-simpl-stats
    Total ticks: 7320

このエラーはGHCの単純化器(simplifier)が無限ループを起こしたときに発生する*3。問題は、pgcd規則が適合して、左辺から右辺への書き換えが起こった後、すぐに右辺のpgcdの呼び出しがまた規則に適合することだ。これはqが0かどうかにかかわりなく起こるので、再帰が止まらない。

この問題に見覚えがある人も居るはずだ。正格な言語でifを関数として実装した場合に、両方の分岐が評価されてしまう問題に似ているのだ。こちらの問題の良く知られた解決策は、評価されて欲しくない部分をラムダで包むことだが、ここではそれだけでは足りない。GHCの単純化器は、ラムダの中であっても簡約できそうなところは何でも簡約するからだ。それでも基本的な考え方は流用できて、要するにifの分岐部分を、「評価が進まない形」にして渡してやれば良い。次の関数を定義しておく。

data Key = Key

delay :: Key -> a -> a
delay _ x = x
{-# NOINLINE delay #-}

{-# RULES "delay/Key" forall x. delay Key x = x #-}

delayは第二引数をそのまま返す関数だが、その事実はNOINLINEプラグマによって隠蔽されるので単純化器が知ることができない*4。例外は、delayの第一引数が具体的にKeyだと分かっている場合であり、このときは規則"delay/Key"が発動して実質的にdelayがインライン展開される。これを使って、(pgcd y q :: Int)という式を次のように書き換える。

(\k -> delay k pgcd y q) :: Key -> Int

この形は、単純化器によって簡約され得る項を一つも含まないので、規則の右辺に持ってきても安全だ。一般には、一つの式の中に簡約を防ぎたいところが複数箇所あっても良い。たとえば(f 3 + g 4)は(\k -> delay k f 3 + delay k g 4)のように変形できる。

このようにして作った「評価の止まった」式を再び解凍して簡約を再開させる関数を用意する。

force :: (Key -> a) -> a
force x = x Key
{-# INLINE force #-}

これによってx内部のdelayの第一引数がKeyに確定するので、"delay/Key"規則によってx内のdelayが除去される。

式の評価を止めることができるようになったので、これを使ってifを実装する。

if_ :: Bool -> (Key -> a) -> (Key -> a) -> a
if_ c x y = if GHC.Exts.lazy c then force x else force y
{-# NOINLINE if_ #-}

{-# RULES
"if_/True" forall x y. if_ True x y = force x
"if_/False" forall x y. if_ False x y = force y
  #-}

"if_/True"と"if_/False"の二つの規則によって、if_の第一引数がTrueかFalseに確定した場合、それに対応する分岐のみが簡約再開される。それ以外の場合、どちらの選択肢も簡約されない*5。

あとは上の規則"pgcd"をif_を使った形に書き直せば良い。

{-# RULES
"pgcd" forall x y. pgcd x y =
    let q = rem x y
    in if_ (q == 0)
        (\_ -> y)
        (\k -> delay k pgcd y q)
  #-}

コード

module M where

import qualified GHC.Exts as GHC

-- | 二つの整数の最大公約数を求める
gcd' :: Int -> Int -> Int
gcd' x 0 = x
gcd' x y = pgcd (abs x) (abs y)
{-# INLINE gcd' #-}

-- | 非負の整数と正の整数の最大公約数を求める
pgcd :: Int -> Int -> Int
pgcd x y
  | q == 0 = y
  | otherwise = pgcd y q
  where
    q = mod x y
{-# NOINLINE pgcd #-}

data Key = Key

delay :: Key -> a -> a
delay _ x = x
{-# NOINLINE delay #-}

{-# RULES "delay/Key" forall x. delay Key x = x #-}

force :: (Key -> a) -> a
force x = x Key
{-# INLINE force #-}

{-# RULES
"pgcd" forall x y. pgcd x y =
    let q = rem x y
    in if_ (q == 0)
        (\_ -> y)
        (\k -> delay k pgcd y q)
  #-}

if_ :: Bool -> (Key -> a) -> (Key -> a) -> a
if_ c x y = if GHC.lazy c then force x else force y
{-# NOINLINE if_ #-}

{-# RULES
"if_/True" forall x y. if_ True x y = force x
"if_/False" forall x y. if_ False x y = force y
  #-}

test :: Int
test = gcd' 120 84

これをコンパイルしてみる。*6

% ghc rulegcd.hs -O -ddump-simpl -fforce-recomp | grep 'M.test' 
M.test :: GHC.Types.Int
M.test = GHC.Types.I# 12

やったね！

*1:将来sumが融合変換の対象になれば、この例はコンパイル時に評価されるようになるだろう

*2:modの代わりにコンパイル時に評価しやすいremを使っている。引数が非負なので結果は同じ

*3:他の状況でも発生するが

*4:実用する場合は、delayの呼び出しが残っても性能に悪影響がないようにNOINLINE[0]とするのが良いかもしれない。この場合後述のRULESにはすべて[~0]を付けること

*5:if_の定義の中でGHC.Exts.lazyを呼んでいるのはやっつけだが、これがないと(if_ (case e of { p -> True; _ -> False }) x y)が(case e of p -> if_ True x y; _ -> if_ False x y)に書き換えられてしまい、結果としてif_の第一引数が確定していないのに"if_/True"と"if_/False"が発動することになる。

*6:このままだと、pgcdの定義の右辺に出てくるpgcdの呼び出しまでもが規則によって書き換えられてしまうので、もうちょっと工夫の余地がある