GHCのタイプチェックプラグインについて、すこし

はじめに

ざっと書いたので、まだ雑です。後ほど訂正などする可能性があります。「まちがい」などは、お手やわらかにご指摘ください。

動機

GHCの型レベル自然数は、かっこいい。リストの長さを型に持たせられるなんて、すごすぎる。でも、残念なところもある。

data List :: Nat -> * -> * where
        Nil :: List 0 a
        (:.) :: a -> List ln a -> List (ln + 1) a

tail_ :: List (n + 1) a -> List n a
tail_ (_ :. xs) = xs

何の問題もない定義のように思える。だけど、これをやると、つぎのようなエラーになる。

error:

  ・Could not deduce: ln ~ n
    from the context: (n + 1) ~ (ln + 1)
      bound by a pattern with constructor:
      .
      .
      .

x :. xsはList (n + 1) a型である。でこれは(:.)を見るとList (ln + 1) aになっているので、n + 1 ~ ln + 1であることはわかる。でも結果の型であるList ln aがList n aとおなじだというこのには納得いかない、とコンパイラが文句を言っている。

つまり、コンパイラはn + 1 == ln + 1からn == lnを導くことができない。

このあたりを解決するのにGHCのタイプチェックプラグインが使える。

実用的には

型レベル自然数について、いろいろな演算を可能にするプラグインを作ってくれている人がいる。実用的には、これを使えばいいと思う。

Hackage: ghc-typelits-natnormalise

上記について、以下の記事で学びました。ありがとうございます。

型レベル自然数(つよい)

何をするか

かんたんなタイプチェックプラグインを作ってみる。

対象読者

Haskellでのプログラミングには慣れている
型レベル自然数についても知っている
GHCの中身に興味がある

また、説明にはStackを使う。

ソースコード

例題のソースコードは下記のあたりに置いてあります。

GitHub: YoshikuniJujo/test_haskell/tribial/qiita/try-typecheck-plugin

なにもしないタイプチェックプラグインを、とりあえず動かす

GHCにはコンパイルパスの途中に処理をつっこむためのプラグインという仕組みがある。これまではCore言語とタイプチェックあたりだけだったのが、最近になって、もっとHaskellのソースコードに近いところにも、プラグインがさしこめるようになった。そのあたりは、つぎの記事にくわしい。

GHC Source Plugin 作ってみた

わかりやすい記事を、ありがとうございます。

GHCの「プラグイン」の仕組みはよくできている。pluginという変数を公開したHaskellのモジュールをプラグインとして使うことができる。その、pluginという変数を束縛する値に、いろいろな種類のプラグインをほうりこめばいい。試してみよう。

% stack new hello-typecheck-plugin
% cd hello-typecheck-plugin
% mkdir src/HelloTypecheckPlugin

まずは、package.yamlのdependenciesにghcを追加する。

dependencies:
- base >= 4.7 && < 5
- ghc

つぎのようなモジュールHelloTypecheckPlugin.Pluginを作成する。

src/HelloTypecheckPlugin/Plugin.hs

{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wall -fno-warn-tabs #-}

module HelloTypecheckPlugin.Plugin where

import GhcPlugins
import TcPluginM
import TcRnTypes

plugin :: Plugin
plugin = defaultPlugin { tcPlugin = const . Just $ TcPlugin {
        tcPluginInit = return (),
        tcPluginSolve = const solveHello,
        tcPluginStop = const $ return () } }

solveHello :: [Ct] -> [Ct] -> [Ct] -> TcPluginM TcPluginResult
solveHello _ _ [] = do
        tcPluginTrace "!HelloTypecheckPlugin.Plugin:" ""
        return $ TcPluginOk [] []
solveHello gs _ ws = do
        tcPluginTrace "!HelloTypecheckPlugin.Plugin:" ""
        tcPluginTrace "Given: " $ ppr gs
        tcPluginTrace "Wanted: " $ ppr ws
        return $ TcPluginOk [] []

コードを解説するまえに、とりあえず使ってみよう。

src/HelloTypecheckPlugin/Test.hs

{-# OPTIONS_GHC -Wall -fno-warn-tabs -fplugin=HelloTypecheckPlugin.Plugin #} 

module HelloTypecheckPlugin.Test where

some :: Int -> Char
some = id

「型チェックに失敗」しないとプラグインは呼ばれないようなので、とりあえずダミーの「型チェックに失敗する値」をいれてある。つぎのように、GHCを呼んでみよう。

% stack build --force-dirty --ghc-options=-ddump-tc-trace 2>&1 | grep -A 15 '!HelloTypecheckPlugin.Plugin'
!HelloTypecheckPlugin.Plugin:
Given: []
Wanted:
  [[WD] hole{co_a45H} {0}:: Int
                            ghc-prim-0.5.3:GHC.Prim.~# Bool (CIrredCan(insol)))]
solveSimpleWanteds end }
.
.
.

Wantedのところに、いろいろと書いてあるけれど、「型チェックを成功させるためには、Int ~ Boolが必要」ということだ。

コードの解説

pluginという変数を束縛すう値がプラグインの本体になる。これはPlugin型の値であり、定義するときにはdefaultPluginのフィールドを置き換えてやればいい。今回はタイプチェックプラグインなので、置き換えるのはフィールドtcPluginの値だ。フィールドtcPluginに置く値の型は、[CommandLineOption] -> Maybe TcPlugin型の値だ。CmmandLineOptionによって、プラグインにオプションがあたえられるが、ここでは説明しない。で、TcPlugin型の値を作ればいいわけだけど、この値は3つのフィールドを持つ。

tcPluginInit
tcPluginSolve
tcPluginStop

tcPluginInitでは前処理をすることができ、その結果として値をtcPluginSolveやtcPluginStopにわたすことができる。今回は説明しない。tcPluginStopでは後処理をすることができるのだと思うが、やはり今回は説明しない。なので、今回見るのはtcPluginSolveだ。フィールドtcPluginSolveにはいる値の型は、つぎのようになる。

s -> [Ct] -> [Ct] -> [Ct] -> TcPluginM TcPluginResult

ここで、型変数sはtcPluginInitからわたされる値なので、ここでは無視する。

なので、つぎの関数だけに注目すれば良いことになる。

solveHello :: [Ct] -> [Ct] -> [Ct] -> TcPluginM TcPluginResult
solveHello _ _ [] = do
        tcPluginTrace "!HelloTypecheckPlugin.Plugin:" ""
        return $ TcPluginOk [] []
solveHello gs _ ws = do
        tcPluginTrace "!HelloTypecheckPlugin.Plugin:" ""
        tcPluginTrace "Given: " $ ppr gs
        tcPluginTrace "Wanted: " $ ppr ws
        return $ TcPluginOk [] []

tcPluginTraceはGHCに-ddump-tc-traceオプションをあたえたときに画面にデバッグ用の情報を表示させるための関数だ。ひとつめの引数はStringであり、ふたつめの引数はSDocだ。ひとつめの引数をラベルとして、ふたつめの引数に「表示させたい値」をpprでSDocにしてあたえる。

EvTermについて(この部分は読みとばしてください)

あとは、だいたいかんたんな話だ。solveHelloの引数gsは「あたえられた制約」であり、引数wsは「求められる制約」であり、gsの制約を利用してwsの制約を解決できるかどうかを検査し、解決できたならTcPluginOk [] []のうちのはじめのリストにその制約(Ct型の値)をいれてやればいい。そうすれば、その制約は解決されたことになる。ただ、(僕にとって)すこし難しかったのがTcPluginOkの第一引数が[Ct]ではなく[(EvTerm, Ct)]だったことだ。EvTermとはいったい何だろうかと悩んだ。

ここでのEvTerm型の値が何なのかを知るためには、HaskellのCore言語に理論的に対応するSystem FCについて学ぶ必要があった。で、ここですこしSystem FCを部分的に解説する。僕自身も十分に深く理解はできていないので、雰囲気をつかむていどの説明だ。GADTsの例で解説する。つぎのようなコードを考える。

(*) 参考: System F with Type Equality Coercions

data Exp a where
        Zero :: Exp Int
        Succ :: Exp Int -> Exp Int
        Pair :: Exp b -> Exp c -> Exp (b, c)

eval :: Exp a -> a
eval Zero = 0
eval (Succ e) = eval e + 1
eval (Pair x y) = (eval x, eval y)

このようなとき、System FCではZeroは、つぎのように定義される。

Zero :: (a ~ Int) => Exp a

パターンマッチでは、つぎのようになる。

eval :: Exp a -> a
eval e = case e of
        Zero (a ~ Int) -> 0 |> (Int ~ a)

わかりやすく言うと、「Zero :: Exp aをマッチするときにはaをIntに置き換えますよ。で、結果の0の型Intは型aに置き換えますよ」ということだ。で、ZeroはつぎのようにInt ~ Intという型強制とともに使われる。

Zero (Int ~ Int)

で、これがcase文でマッチさせられて、つぎのようになる。

Zero (Int ~ Int) -> 0 |> (Int ~ Int)

このようにして、ZeroはExp IntからExp Intに、0はIntからIntに型を強制される。

パターンマッチによって、(a ~ Int)という型強制を取り出すことができて、それを使うことで(Int ~ a)という型の置き換えができる。このようにすることで、関数evalのなかで型の整合性を保つことができる。それに対してZeroは(Int ~ Int)という型強制を持ってやってくるので、引数にZeroを適用した結果の型は、めでたくIntになるといった感じ。

ここで、0は(Int ~ a)という型強制によって、型Intから型aに置き換えられているのだけど、この(Int ~ a)の部分を指定してやるのが、EvTermということになる。

EvTermを見てみる(この部分も読みとばしてください)

実際に指定されたEvTermが使われているところを見てみよう。

モジュールHelloTypecheckPlugin.Testの値someを消して、つぎのような定義を書く。

src/HelloTypecheckPlugin/Test.hs

{-# LANGUAGE GADTs #-}
{-# OPTIONS_GHC -Wall -fno-warn-tabs -fplugin=HelloTypecheckPlugin.Plugin #-}

module HelloTypecheckPlugin.Test where

data Foo a where
        FooZero :: Foo a

eval :: Foo a -> a
eval FooZero = (0 :: Int)

もちろん、これもコンパイルは通らない。

% stack build --force-dirty --ghc-options=-ddump-trace 2>&1 | grep -A 15 '!HelloTypecheckPlugin.Plugin'
!HelloTypecheckPlugin.Plugin:
.
.
.
!HelloTypecheckPlugin.Plugin:
Given: []
Wanted:
  [[WD] hole{co_a3Mp} {1}:: a_a3Me[sk:1]
                            ghc-prim-0.5.3:GHC.Prim.~# Int (CNonCanonical)]
evalSimpleWanteds end }
.
.
.

(a ~ Int)という制約の解決が必要とのことだ。モジュールHelloTypecheckPlugin.Pluginにつぎのようにコードを追加する。

src/HelloTypecheckPlugin/Plugin.hs

{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wall -fno-warn-tabs #-}

module HelloTypecheckPlugin.Plugin where

import Data.Maybe

import GhcPlugins
import TcPluginM
import TcRnTypes

import TyCoRep
import TcEvidence

plugin :: Plugin
plugin = defaultPlugin { tcPlugin = const . Just $ TcPlugin {
        tcPluginInit = return (),
        tcPluginSolve = const solveHello,
        tcPluginStop = const $ return () } }

solveHello :: [Ct] -> [Ct] -> [Ct] -> TcPluginM TcPluginResult
solveHello _ _ [] = do
        tcPluginTrace "!HelloTypecheckPlugin.Plugin:" ""
        return $ TcPluginOk [] []
solveHello gs _ ws = do
        tcPluginTrace "!HelloTypecheckPlugin.Plugin:" ""
        tcPluginTrace "Given:" $ ppr gs
        tcPluginTrace "Wanted: " $ ppr ws
        return $ TcPluginOk (catMaybes $ mkEvTerm <$> ws) []

mkEvTerm :: Ct -> Maybe (EvTerm, Ct)
mkEvTerm ct@(CNonCanonical (CtWanted (TyConApp _ [_, _, a, i]) _ _ _)) = Just (
        EvExpr . Coercion
                $ mkUnivCo (PluginProv "!HelloTypecheckPlugin.Plugin") Nominal a i,
        ct)
mkEvTerm _ = Nothing

関数mkEvTermはCt型の値から型aと型Intの値を取り出して、EvExpr . Coercion $ mkUnivCo ...でUnivCoという型強制の値(EvTerm型の値)を作成してい、Ct型の値にくっつけている。そのEvTerm型の値をくっつけたものをTcPluginOkの第一引数のリストにいれてやることで、この型制約は満たされたことになる。ここでは、とくに型のなかみをチェックしていないので、わりと無差別にどんな型のあいだでも、型強制が生じてしまうことになる。

これで、とりあえずコンパイルが通るようになる。

% stack build

さて、EvTerm型の値はどうなったか。それを見てみよう。CoreからCoreへの変換をおこなうプラグインの仕組みを流用して、タイプチェックプラグインの結果を見てみよう。

つぎのような関数を定義する。

src/HelloTypecheckPlugin/Plugin.hs

showCore :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
showCore _ = return . (corePlugin :)

corePlugin :: CoreToDo
corePlugin = CoreDoPluginPass "!HelloTypecheckPlugin.Plugin" $ \mg -> do
        putMsg . ppr $ mg_binds mg
        return mg

こうしたうえで、defaultPluginのフィールドを書き換える。

src/HelloTypecheckPlugin/Plugin.hs

plugin :: Plugin
plugin = defaultPlugin {
        installCoreToDos = showCore,
        tcPlugin = const . Just $ TcPlugin {
                ... } }

ビルドしてみる。

% stack build
.
.
.
 [LclIdX]
 eval
   = \ (@ a) (ds_dsZQ :: Foo a) ->
       case ds_dsZQ of { FooZero ->
       (ghc-prim-0.5.3:GHC.Types.I# 0#)
       `cast` (Univ(representational plugin "!HelloTypecheckPlugin.Plugin"
                      :: Int, a)
               :: Coercible Int a)
       }
.
.
.

ずらずらっと流れた出力のなかにevalの定義があり、そのなかに、つぎのような部分がある。

`cast` (Univ(representational plugin "!HelloTypecheckPlugin.Plugin" :: Int, a) :: Coercible Int a)

この部分がEvTermとして設定したところだ。

1たしたものどうし等しければ等しい(途中読みとばした人はここから続きを読んでください)

さて、EvTermとはという話を長々としてきた。だけど、今回やることにEvTermはそんなに関係がない。まえにもみたが、つぎのようなコードを考える。

% stack new typecheck-plusone
% cd typecheck-plusone
% mkdir src/TypecheckPlusone

つぎのようなコードを書く。

src/TypecheckPlusone/Test.hs

{-# LANGUAGE GADTs, DataKinds, TypeOperators, KindSignatures #-}
{-# LANGUAGE StandaloneDerivings #-}
{-# OPTIONS_GHC -Wall -fno-warn-tabs #-}

module TypecheckPlusone.Test where

import GHC.TypeLits

infixr 5 :.

data List :: Nat -> * -> * where
        Nil :: LIst 0 a
        (:.) :: a -> List ln a -> LIst (ln + 1) a

deriving instance Show a => Show (List ln a)

tail_ :: List (n + 1) a -> List n a
tail_ (_ :. xs) = xs

これは、うえでみたようにコンパイルエラーになる。

% stack build
.
.
.
    ・Could not deduce: ln ~ n
      from the context: (n + 1) ~ (ln + 1)
.
.
.

つまり、n + 1 == ln + 1なのはわかるけど、n == lnなのはわからないよ、ということだ。なので、nとlnの等値性を(n + 1)と(ln + 1)の等値性に置き換えられればいい。

つぎのようなコードを書く。

src/TypecheckPlusone/Plugin.hs

{-# LANGUAGE TupleSections #-}
{-# OPTIONS_GHC -Wall -fno-warn-tabs #-}

module TypecheckPlusone.Plugin where

import Data.Maybe

import GhcPlugins
import TcPluginM
import TcRnTypes
import TyCoRep
import TcEvidence
import TcTypeNats

plugin :: Plugin
plugin = defaultPlugin { tcPlugin = const . Just $ TcPlugin {
        tcPluginInit = return (),
        tcPluginSolve = const $ solvePlusOnePlugin,
        tcPluginStop = const $ return () } }

solvePlusOnePlugin :: [Ct] -> [Ct] -> [Ct] -> TcPluginM TcPluginResult
solvePlusOnePlugin gs _ ws = do
        (oks, news) <- unzip . catMaybes <$> solvePlusOne gs `mapM` ws
        return $ TcPluginOk oks news

solvePlusOne :: [Ct] -> Ct -> TcPluginM (Maybe ((EvTerm, Ct), Ct))
solvePlusOne gs w
        | Just (t1, t2) <- getTypes w = do
                if any (`isPlusOneOf` t1) gts && any (`isPlusOneOf` t2) gts
                then do nct <- makeCt w (plusOne t1) (plusOne t2)
                        return $ Just ((makeEvTerm t1 t2, w), nct)
                else return Nothing
        | otherwise = return Nothing
        where gts = allTypes gs

getTypes :: Ct -> Maybe (Type, Type)
getTypes ct = case classifyPredType . ctEvPred $ ctEvidence ct of
        EqPred NomEq t1 t2 -> Just (t1, t2); _ -> Nothing

allTypes :: [Ct] -> [Type]
allTypes cts =
        concat . catMaybes $ ((\(t1, t2) -> [t1, t2]) <$>) . getTypes <$> cts

plusOne :: Type -> Type
plusOne t = TyConApp typeNatAddTyCon [t, LitTy $ NumTyLit 1]

isPlusOneOf :: Type -> Type -> Bool
TyConApp tc [TyVarTy t1, LitTy (NumTyLit 1)] `isPlusOneOf` TyVarTy t =
        tc == typeNatAddTyCon && t1 == t
_ `isPlusOneOf` _ = False

makeCt :: Ct -> Type -> Type -> TcPluginM Ct
makeCt ct t1 t2 = do
        hole <- newCoercionHole pt
        return . mkNonCanonical $ CtWanted pt (HoleDest hole) WDeriv (ctLoc ct)
        where pt = mkPrimEqPred t1 t2

makeEvTerm :: Type -> Type -> EvTerm
makeEvTerm t1 t2 = EvExpr . Coercion
        $ mkUnivCo (PluginProv "TypecheckPlusone.Plugin") Nominal t1 t2

とりあえず、これを書いたら、モジュールTypecheckPlusone.TestのOPTIONS_GHCを書き換える。

src/TypecheckPlusone/Test.hs

{-# OPTIONS_GHC -Wall -fno-warn-tabs -fplugin=TypecheckPlusone.Plugin #-}

これで、コンパイルが通るようになる。

コードの説明

大きな流れとしては、つぎのとおりだ。

受け取った制約(gs)があり
必要な制約(ws)がある
wsがn ~ mのかたちの制約であったなら
- gsにn + 1とm + 1が存在することを確認し
- 存在していたら、この制約は解決したものとし
- 新たに(n + 1) ~ (m + 1)という制約を追加する

関数solvePlusOnePluginは関数solvePlusOneで、「単一の必要な制約」に対して行っている処理を「複数の必要な制約」に対して行うようにしているだけだ。本質的な処理は関数solvePlusOneが担っている。

solvePlusOneではgetTypesでふたつの型を取り出し、それぞれに対応する(+ 1)された型がgsに含まれるかをチェックして、含まれていれば、新しい制約を作成し、満たされた制約やEvTermとともにタプルを組み立てている。

getTypesでは、関数classifyPredType, ctEvPred, ctEvidenceなどによって、含まれる(t1 ~ t2)というかたちを取り出している。allTypesはgetTypesによって、複数のCtから取り出された型を、均一なリストにまとめる。

makeCtはnewCoercionHoleで型強制を入れるための穴を作ってやり、それを使って新しい「必要な制約」を生成している。このとき、CtLoc(たぶん、ソースコード上での位置)が必要になるが、これは第1引数のCt型の値のものを、そのまま使っている。

makeEvTermはEvTermを作っている。EvTermについては上記「読みとばすべき部分」を参照のこと。

まとめ

タイプチェックプラグインを作ってみた。GHCのプラグインはよくできていて、ふつうのモジュールをつくり、pluginという変数を束縛する値を公開してやれば、それをそのままプラグインとして使うことができる。

今回作ってみたのは、n + 1 == m + 1からn == mを引き出すことのできる型チェックプラグインだ。