コンテキストつきmapとStateモナドのこと


[a]があって、それにある関数 f :: a -> b -> (c, b) をマップして[c]を得たい。ここで、bはマップ開始時の初期値は決まっているけど、それから先はマップの各回によって新しいbの値が生成される。

こんなようなコードを書く必要があるときがありました。きっとStateモナドはこうゆうときに使うのだろうなーという気はするのですが、いかんせん使ったことがない、サンプルを見てもいまいち理解できない、で、こんなコードを書いてしまいます。

compute :: Int -> Int -> (Int, Int)
compute a c = (a + c, c + 1)

mapWithContext :: (a->c->(b,c)) -> c -> [a] -> [b]
mapWithContext f c  = 
mapWithContext f c (x:xs) = (x':mapWithContext f c' xs)
	where
		(x', c') = f c x

test :: IO()
test = putStrLn $ show $ 
        mapWithContext compute 1 [1..5]

こうゆうコードを書いたことは今まで1回だけではなくて、常々思っていたのは、1)再帰を自分で書くのはできるだけ避けるべきだ。(mapWithContext) 2)こうゆう振る舞いをするコードを書く必要のある状況というのはそんなに特殊じゃないはずだ...ということでした。

今回、IOモナドが落ち着いたので、Stateモナドを使ってこのコードを書き換えてみることにしました。サンプルを色々あたってみたのですが、単純すぎて話にならないものもあったり、ほかにはツリー構造の各要素に番号を振っていくようなのもあったのですが、気になったのは、そのサンプルが再帰を自分でやっていることでした。Stateモナドを使っても、結局再帰の部分を自分で書かなくてはならないとしたら少し悲しいものがあります。ツリー構造をトラバースするにしても、トラバースのコードと各要素でやりたい仕事を行うコードは別の関数に分けることが望ましいですよね…

Stateモナドの定義は

State st a = State {runState :: (st -> (a, st))}

ということで、モナドの中にステート値をパラメータにとり、値とステート値のタプルを返す関数を内包しているようです。先ほどのコードに出てくるcomputeのパラメータcをステート値と見立てれば、computeは以下のように書き換えられます。

compute :: Int -> State Int Int
compute a = do
 c <- get
 put (c + 1)
 return (a + c)

getとputはMonadStateクラス関数で、Stateモナドからステート値(st)をとりだしたり、書き換えたりすることができます...でも、書き換えって副作用じゃんと思ってしまうのですが、どうやらそうではないんですね…
do記法は実は以下のような表現のシンタックスシュガーらしいです。

compute :: Int -> State Int Int
compute a = 
  get >>= \c ->
 put (c + 1) >>= \_ ->
 return (a + c)

つまり、do表現の各行はbindでつながれた一本の式だったんですね…
そして、
 (>>=) :: m a -> (a -> m b) -> m b
ですから、つまり、do表現の各行で新しいm b、ここではStateモナドが生成されているということらしいです。

getとputのコードは

get = State $ \s -> (s,s)
put s = State $ \_ -> ((),s)

なので、getはstをStateモナドのaの部分にも持っているStateモナドを返し、putはパラメータsをstとして持ち、aは()なStateモナドを返す...それが次のbindに渡されるとbindの中でaが取り出されてbindの右辺の関数にわたる。つまりbind間の各項を実行するたびに新しいStateモナドができているので、副作用を起こさずにaやstを変化させられる。

ということらしいです。でも、なんだかこの例のcomputeはStateモナドを使う前よりも大きくなっちゃいました...ここはとりあえずほっといて、呼び出し側(mapWithContext)を見てみましょう。computeの型が変わったのでそれの対応をまずします。
さらに、mapWithContextにもステート情報cが渡っているので、この関数もStateモナドを使うことができそうです…
でもcomputeを呼び出すたびに変化していくステートを次のcomputeに渡していくにはどうしたらよいのだろうと、考えてしまいました。

mapWithContext :: (a -> State c b) -> [a] -> State c [b]
mapWithContext f  = 
mapWithContext f lst@(x:xs) = do
	x' <- f x
	xs' <- mapWithContext f xs
	return (x' : xs')
 

なんていう手もありますが、もう一つ考えたのは、単純にfをlstにマップすれば[State c b]ができるから、それをsequenceに渡せばState c [b]になったりしたらかっこいいなーと思ったので、やってみました。

mapWithContext :: (a -> State c b) -> [a] -> State c [b]
mapWithContext f lst = sequence $ map f lst

すると、これが動いちゃうんですね。かっこいい!どうやらリストのトラバースの場合は自分で再帰のコードは書かなくてよいようです。さらに調べてみるとmapM = sequence .mapだそうなので、

mapWithContext :: (a -> State c b) -> [a] -> State c [b]
mapWithContext f lst = mapM f lst

ということで、mapWithContext丸ごといらなくなってしまいました。結果:

compute :: Int -> State Int Int
compute a = get >>= \x -> put (x + 1) >> return (a + x)

test :: IO()
test = putStrLn $ show $ 
        evalState (mapM compute [1..5]) 1

さらにcomputeの中のgetとputが邪魔くさいと思ったので、色々やってみたのですが、ステートモナドは関数を内包しているので、computeを以下のように書き換えても動くことがわかりました。

compute :: Int -> State Int Int
compute a = State $ \x -> (a + x, x + 1)

これでずいぶんさっぱりしました。今回のサンプルではgetやputを使うより、こっちのほうがバランスが取れてる気がします…

ではでは