Template Haskell 入門した

mr_konn さんの できる!Template Haskell (完) を読んだので自分で書いてみたメモ。

動機

最近書いているコードで

data Hoge = Hoge1 | Hoge2 | Hoge3 deriving (Eq, Show)
data Fuga = Fuga1 | Fuga2 deriving (Eq, Show)

みたいなデータに対し、

convertHoge :: String -> Hoge
convertHoge t
  | t == "hoge1" = Hoge1
  | t == "hoge2" = Hoge2
  | t == "hoge3" = Hoge3
  | otherwise    = error "Hoge"

convertFuga :: String -> Fuga
convertFuga t
  | t == "fuga1" = Fuga1
  | t == "fuga2" = Fuga2
  | otherwise    = error "Fuga"

のような変換関数を1つずつ書くのが嫌だったので、 このような関数を自動的に生成できるようなものを書こうとしてました。

これを生成する関数を mkConvertFunc とすると

mkConvertFunc ''Hoge ["hoge1", "hoge2", "hoge3"]

convertHoge が生成される感じ。

単純な read ではないことに注意 (capitalize しただけでない場合もある) 。 ただ、こういう関数書くことはちょくちょくあると思うので どっかにライブラリとして転がっているような気もする。

実装

構文木と TH 版 convert 関数

まず runQ を使って構文木の形を調べる

$ ghci -XTemplateHaskell
> :m Language.Haskell.TH
> runQ [| \t -> | t == "hoge1" = Hoge1; | t == "hoge2" = Hoge2; | t == "hoge3" = Hoge3; | otherwise = error "Hoge" |]
<interactive>:65:15: parse error on input `|'

ふむ、ラムダ式ではパターンマッチのガード | を使えないらしい。 case を使って書きなおし。

> runQ [| \t -> case t of "hoge1" -> Hoge1; "hoge2" -> Hoge2; "hoge3" -> Hoge3; _ -> error "Hoge" |]
LamE [VarP t_0]
  (CaseE (VarE t_0)
    [ Match (LitP (StringL "hoge1")) (NormalB (ConE Main.Hoge1)) []
    , Match (LitP (StringL "hoge2")) (NormalB (ConE Main.Hoge2)) []
    , Match (LitP (StringL "hoge3")) (NormalB (ConE Main.Hoge3)) []
    , Match WildP (NormalB (AppE (VarE GHC.Err.error) (LitE (StringL "Hoge")))) []
    ]
  )

できた。出力は読みやすいように整形してます。

これをほぼそのまま書き下して、 TH 版 convertHoge を書ける。

convertHogeTH :: ExpQ
convertHogeTH = do
  t_0 <- newName "t"
  return $ LamE [VarP t_0]
    (CaseE (VarE t_0)
      [ Match (LitP (StringL "hoge1")) (NormalB (ConE 'Hoge1)) []
      , Match (LitP (StringL "hoge2")) (NormalB (ConE 'Hoge2)) []
      , Match (LitP (StringL "hoge3")) (NormalB (ConE 'Hoge3)) []
      , Match WildP (NormalB (AppE (VarE 'error) (LitE (StringL "Hoge")))) []
      ]
    )

Hoge1 とか error みたいな関数 (値) はクオートすることに注意。

試してみる

> :l Convert.hs
> $(convertHogeTH) "hoge1"
Hoge1
> $(convertHogeTH) "hoge2"
Hoge2
> $(convertHogeTH) "hoge"
*** Exception: Hoge

大丈夫そう。

簡易版 mkConvertFunc' の実装

ではこれに HogeFuga を引数に渡してどちらにも使えるようにしてみる。いきなり目的の

mkConvertFunc ''Hoge ["hoge1", "hoge2", "hoge3"]

を作るのは難しそうだったので、 String -> Hoge へのマッピング用のタプルのリストを引数に渡して

mkConvertFunc' ''Hoge [("hoge1", Hoge1), ("hoge2", Hoge2), ("hoge3", Hoge3)]

を作った。

mkConvertFunc' :: Name -> [(String, Name)] -> ExpQ
mkConvertFunc' d lst = do
  t_0 <- newName "t"
  return $ LamE [VarP t_0]
    (CaseE (VarE t_0)
      $ (map (\(s, t) -> Match (LitP (StringL s)) (NormalB (ConE t)) []) lst)
      ++ [Match WildP (NormalB (AppE (VarE 'error) (LitE (StringL $ show d)))) []])

実際これは難しくないのだけど、型が

mkConvertFunc' :: Name -> [(String, Name)] -> ExpQ

になるのが最初わからなかった。 まあ考えればわかるんだけど、 ExpQ とか ConE ではなく Name 。 ドキュメントちゃんと読めってことですかね。 http://hackage.haskell.org/package/template-haskell

あと、6行目 $ (map ... ) が一見 $( ... ) に見えるが実は関数結合の $ です。 これを書いた後から見なおして早速混乱してしまった。要注意。 TH を書くときは関数結合の $ を使わないほうが良さそう (と、 mr_konn さんの記事にも書いてある) 。

> $(mkConvertFunc ''Hoge [("hoge1", 'Hoge1), ("hoge2", 'Hoge2), ("hoge3", 'Hoge3)]) "hoge1"
Hoge1
> $(mkConvertFunc ''Hoge [("hoge1", 'Hoge1), ("hoge2", 'Hoge2), ("hoge3", 'Hoge3)]) "hoge"
*** Exception: Hoge

mkConvertFunc の実装

さて、では最後に目的の

mkConvertFunc ''Hoge ["hoge1", "hoge2", "hoge3"]

を作る。これには data の宣言

data Hoge = Hoge1 | Hoge2 | Hoge3 deriving (Show, Eq)

から Hoge1 Hoge2 Hoge3 という3つのコンストラクタを取ってくることが必要で、 reify という関数を使う。 reify はコンパイル時にしか使えないので ghci から叩けない。 ソースコード中に

$(do
   info <- reify ''Hoge
   runIO $ print info
   return [])

と書いてみて、 ghci からロードする。

> :l Convert.hs
TyConI (DataD [] Convert.Hoge [] [NormalC Convert.Hoge1 [],NormalC Convert.Hoge2 [],NormalC Convert.Hoge3 []] [])

なるほど。こいつから Hoge1 Hoge2 Hoge3 を取る。

$(do
  ctrs <- (\(TyConI (DataD [] _ [] x [])) -> map (\(NormalC name []) -> name) x)
          <$> reify ''Hoge
  runIO $ print $ ctrs
  return [])
> :r
[Convert.Hoge1,Convert.Hoge2,Convert.Hoge3]

ちょっと見にくいけどこれでよさそう。 最後にこれを使って mkConvertFunc を定義して終わり。

mkConvertFunc :: Name -> [String] -> ExpQ
mkConvertFunc d strs = do
  ctrs <- (\(TyConI (DataD [] _ [] x [])) -> map (\(NormalC name []) -> name) x)
          <$> reify d
  v <- newName "x"
  return $ LamE [VarP v]
    (CaseE (VarE v)
       $ (map (\(s,t) -> Match (LitP (StringL s)) (NormalB (ConE t)) []) $ zip strs ctrs)
       ++ [Match WildP (NormalB (AppE (VarE 'error) (LitE (StringL $ show d)))) []])
> :r
> $(mkConvertFunc ''Hoge ["hoge1", "hoge2", "hoge3"]) "hoge1"
Hoge1
> $(mkConvertFunc ''Hoge ["hoge1", "hoge2", "hoge3"]) "hoge2"
Hoge2
> $(mkConvertFunc ''Hoge ["hoge1", "hoge2", "hoge3"]) "hoge3"
Hoge3
> $(mkConvertFunc ''Hoge ["hoge1", "hoge2", "hoge3"]) "hoge4"
*** Exception: Convert.Hoge
> $(mkConvertFunc ''Fuga ["fuga1", "fuga2"]) "fuga1"
Fuga1
> $(mkConvertFunc ''Fuga ["fuga1", "fuga2"]) "fuga2"
Fuga2
> $(mkConvertFunc ''Fuga ["fuga1", "fuga2"]) "fuga3"
*** Exception: Convert.Fuga

良さそう。

一応まとめたものを Gist