MatlabのDSL in Haskell

今回例として扱うのは、Matlab（行列計算言語・統合環境）の行列計算を表現するようなDSL。最小構成として、２次元行列の加減乗算と要素の取り出しを表現したい。（Matlabの文法の簡単な説明はこれなどを参照。 http://www.math.meiji.ac.jp/~mk/labo/text/matlab/node4.html）これだけの小さなものでも、

• 変数の管理
• let-sharing（別の記事で記述予定）

といった、コード生成のエッセンスが含まれている。

まず、構文木を表すデータ型Exp rはこのような感じになる。

import Data.Text (Text)
import qualified Data.Text as T

type Ident = Text

-- 2D matrix
data Mat2D r = Mat Ident | MatConst [[r]]

class Scalar a
instance Scalar Int
instance Scalar Integer
instance Scalar Double

data Exp r where
	Const :: r -> Exp r
	Add :: Exp r -> Exp r -> Exp r
	Sub :: Exp r -> Exp r -> Exp r
	Mul :: Exp r -> Exp r -> Exp r
	Abs :: (Scalar r) => Exp r -> Exp r
	Signum :: (Scalar r) => Exp r -> Exp r
	Index2D :: (Num r) => Exp Int -> Exp Int -> Exp (Mat2D r) -> Exp r

この構文木を使って、たとえば、reify :: Exp r -> Textという関数を定義して、

> reify $ Const (MatConst [[1,2,3],[4,5,6]])
[[1 2 3];[4 5 6]]
> reify $ Index2D 1 2 (MatConst [[1,2,3],[4,5,6]])
m = [[1 2 3];[4 5 6]];
m(1,2)

といった具合のコード生成をしたい。「コード生成できるもの」を表すReifiableクラスというものを作って、reify :: a -> TextをReifiableクラスのメソッドとする。ここでReifiableのインスタンスとなるのは、Exp r, Mat2D r, Int, Integer, Doubleである。例として以下にExp rとMat2D rのインスタンス定義を示す。ここまでの全ソースコードと実行例はここにある。

instance Reifiable (Exp r) where
	reify (Const exp) = reify exp
	reify (Add e1 e2) = T.concat [reify e1,"+",reify e2]
	reify (Index2D i j mat) = T.concat [reify mat,"(",reify i,",",reify j,")"]

instance (Show r) => Reifiable (Mat2D r) where
	reify (Mat ident) = ident
	reify (MatConst xss) = T.pack $ "[" ++ intercalate ";" (map (intercalate " " . map show) xss) ++ "]"

ちなみに、ここでは行列は二次元のみであるが、行列を２次元以外にも対応させて型安全にするには、ここ（http://www.haskell.org/haskellwiki/GHC/Type_families）にあるように、data family Mat dim rなどとして、dimにD1, D2, D3といった型のタグを付けてやれば良いと思われる。（ただしそのようなことをやると、型クラスをつけるときにだいぶ複雑になる。）

ASTは単なるHaskellの値であるので、mapやfoldなどを使ってASTを構築することもできる。

*Main> reify $ foldl1 (+) $ replicate 5 $ Const (MatConst [[1,2,3],[4,5,6]])
"[1 2 3;4 5 6]+[1 2 3;4 5 6]+[1 2 3;4 5 6]+[1 2 3;4 5 6]+[1 2 3;4 5 6]"

なかなか良い感じなのだが、実はこの実装では問題があり、Index2Dによる要素の取り出しがきちんと出力できない。
Index2D 1 2 (MatConst [[1,2,3],[4,5,6]])をreifyすると、[[1,2,3],[4,5,6]](1,2)となるのだが、Matlabのいけてない言語仕様により、このようには書けず、定数行列を一回変数に代入する必要がある！

モナドを利用した構造化

• そういった、コード出力時に単純変換以上の何らかの操作をしたい場合には、状態管理を出力のロジックに組み入れる必要がある。→ Stateモナドを使う。
• また、コードの出力も、単にTextを返すのではなく、ASTをたどっていく過程の任意の段階で出力をしたい。→ Writerモナドを使う。

State、Writerの両方を使いたいので、モナド変換子で２つのモナドを積み上げるのが基本のようだが、今回はRWSTというモナド変換子（Hackage、Shelarcyさんによる解説）を使う。RWSモナドはReader, Writer, Stateのすべてが組み込まれており、しかもMonadState(http://hackage.haskell.org/packages/archive/mtl/2.0.1.0/doc/html/Control-Monad-State-Class.html)などの型クラスのインスタンスになっているので、get/put, tellなどがliftせずに使えるようになっている便利なモナドである。IOも、後ほどのlet-sharingの件で必要となるので入れている。

type W a = RWST () Text CompilerState IO a

最初この巨大な型を見ると面食らうが、その意味するところは、型パラメータの左から順に、読み込む（Reader）ものは無く（()）、Textを書き出すことができ、CompilerStateという状態を書き換えることができ、IOもliftIOを介して使える、というものである。書きだすもの（Text）と状態（CompilerState）は型が決まっており、get,put,tellなどの関数を使って、モナドの裏側で暗黙に操作される（これらの関数はみなW a型の値を返す。）。 型コンストラクタの最後に付くa型はdo記法の中で自由に替えられ、モナドのユーザーが明示的に扱うものであるというのは、IOなど他のモナドと同じである。

最初の、reifyがTextを返すバージョンの抜粋

class Reifiable a where
	reify :: a -> Text

instance Reifiable (Exp r) where
	reify (Const exp) = reify exp
	reify (Add e1 e2) = T.concat [reify e1,"+",reify e2]
	reify (Index2D i j mat) = T.concat [reify mat,"(",reify i,",",reify j,")"]

gen :: Exp r -> IO ()
gen = T.putStrLn . reify

reifyがモナドを返すように書き換えたバージョンの抜粋。

type W a = RWST () Text CompilerState IO a
class Reifiable a where
	reify :: a -> W Text

instance Reifiable (Exp r) where
	reify (Const exp) = reify exp
	reify (Add e1 e2) = do
		c1 <- reify e1
		c2 <- reify e2
		return $ T.concat [c1,"+",c2]
	reify (Index2D i j mat) = do
		ci <- reify i
		cj <- reify j
		m <- reify mat
		return $ T.concat [m,"(",ci,",",cj,")"]

data CompilerState = CompilerState { variableCount :: Int }

gen :: Exp r -> IO ()
gen e = do
	(txt,_,_) <- runRWST (reify e) () (CompilerState 0)
	T.putStrLn txt

まだ機能は古いバージョンと全く同じ。runRWSTを使ってモナドを走らせると、IO (Text, CompilerState,Text)を返すので、その結果を取り出せる。

先に述べたように、定数行列から要素を取り出すとき、いったん行列を変数に代入して、その変数から要素取り出しをしたい。
どのように設計すればよいかというと、

• 以前と同様、reifyが返すのはその値を参照するのに必要なコード。つまり代入した変数の名前。
• それに加えて、MatConstに対するreifyの定義の中で、tellを使って、値を準備するコード（つまり新しい変数への代入のコード）をWriterモナドに「出力」する。

instance (Show r) => Reifiable (Mat2D r) where
	reify (Mat ident) = return ident
	reify (MatConst xss) = do
		let code =  T.pack $ "[" ++ intercalate ";" (map (intercalate " " . map show) xss) ++ "]"
		var <- newName
		tell $ T.concat [var," = ",code,";\n"]
		return var

ここで、「出力」と鍵括弧にくくったのは、putStrLnのようにIOに出力するのではなく、あくまでWriterモナドをつかっているので、最後にモナドを走らせて「出力」を取り出す必要があるからである。newNameというのは、Wモナドが持つ状態（WモナドはMonadStateのインスタンスである）を利用して一意な名前を作る関数で、次のように定義できる。

data CompilerState = CompilerState { variableCount :: Int }

newName :: W Text
newName = do
	CompilerState count <- get
	put $ CompilerState (count + 1)
	return $ T.pack $ "v" ++ show (count+1)

最後に、Writerモナドに出力した内容と、reify自体の返り値をまとめてコード出力の結果とする。

gen :: Exp r -> IO ()
gen e = do
	(txt,_,assigned) <- runRWST (reify e) () def
	T.putStrLn $ T.append assigned txt

実行結果は以下のようになる。

*Main> let m1 = Const (MatConst [[1,2,3],[4,5,6]]) :: (Exp (Mat2D Int))
*Main> let m2 = Const (MatConst [[10,5,3],[1,4,2]]) :: (Exp (Mat2D Int))
*Main> gen $ Index2D (Index2D 1 2 m1) 1 m2
v1 = [1 2 3;4 5 6];
v2 = [10 5 3;1 4 2];
v2(v1(1,2),1)

良い感じに変数への代入ができている！　コードをここに挙げておく。

変数の重複の排除

ここで、以下の様な式を考えてみる。

m1 :: Exp (Mat2D Int)
m1 = Const (MatConst [[1,2,3],[4,5,6]])

test4 :: Exp (Mat2D Int)
test4 = m1 + m1

これをreifyすると、

*Main> gen test4
v1 = [1 2 3;4 5 6];
v2 = [1 2 3;4 5 6];
v1+v2

となり、同じ内容であるはずのm1がv1,v2と複数回生成されている。これくらいの小さなデータならいいが、もしこのm1が重い計算の結果得られる値の場合このような重複は避けたい。なぜこういうことが起こるかというと、test4が簡約されるときに、m1 + m1という式の両側のm1が同じがどうかというのを構わずにHaskellのコンパイラが展開してしまうからである。この論文の3.1節に説明されている。これの重複を避けることには、Let-sharingという名前が付いている。一言で言えば、System.Mem.StableNameモジュールのmakeStableNameとhashStableNameという関数を使って、項の同一性判定をしてやれば良く、さほど難しくはなかった。別の記事で実装について記す予定。

参考になった記事

• http://d.hatena.ne.jp/keigoi/20111206/haskell_tagless_dsl
• http://blog.kzfmix.com/entry/1340796741

文献を調査中なので、参考になるものがあったら教えてほしいです。