Stateモナドを絵を描いて理解してみた

絵を描きながらStateモナドを考えてみたら意外とうまくいったので紹介します。
Stateモナドには「表の値」と「裏の値=状態」という2つの値(型)が登場します。
Stateモナドの定義です。(s -> (a,s))で、sが状態の型、aが表の値の型です。

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

The State monad

図の方ですが、矢印は関数です。原則として、1入力1出力の一本の矢印で表しますが、今回は戻り値が2値のタプルなので2出力として書きました。
丸四角で囲ってあるのは、関数をオブジェクト*1として扱うことを意図しています。
これに対して、丸四角で囲っていない矢印は関数適用を表すことにします。
一番外側の大カッコはStateモナドのデータ構築子を表します。これでラムダを囲むとStateモナドの値になります。
runStateは、Stateの値から関数を取り出すという役割だけなので、ここでは特に図にしていません。
観察すると、Stateの実体は単なる関数だということがわかります。

次に return と >>=(bind)です。

instance Monad (State s) where 
    return a        = State $ \s -> (a,s)
    (State x) >>= f = State $ \s -> let (v,s') = x s in runState (f v) s'

The State monad

右辺を図にしていきます。
まずreturn a。

上の矢印が表の値側で、下が状態側です。矢印は恒等関数(id)です。return a は、値は常にa、状態は受け取ったものをそのまま出します。

(State x) >>= f (bind)です。

新しい記法として点線が出てきましたが、中の矢印をコピーしているイメージです。ということは、これはrunStateです。Stateモナドの中から実体の関数を取り出します。
図では、取り出した関数をs'に適用しています。コードの runState (f v) s' の部分です。
式で見るよりすっきりした感じですが、まだごちゃごちゃしています。
わざわざモナドの中からコピーして取り出して適用しなくても、そのままつないでしまっていいんじゃないか、ということで次の図になりました。

丸四角はラムダだとしてしまったので、入出力の矢印が丸四角を飛び出したら関数適用ということにします。
矢印に無駄がなくなりましたが、曲がりくねっているのが不満です。
よく見ると、>>=の図で今いじくった部分は全部関数 f の担当のところです。大雑把に、そこら辺をもう全部 f としてしまいます。乱暴な感じもしますが、こう考えると納得できます。
f は State モナドの値を返す。State モナドは実質的にラムダなので、f はカリー化をされている、カリー化を陽に書いていると思える。したがって、 f は実質的に2変数関数です。
こう考えると次の図になります。

f と書いてある四角はブラックボックスを表します。ごちゃごちゃしてたところを全部四角の中に押し込んだので、2入力2出力の関数というだけのすっきりした図になりました。
これをみると f は、表の値と状態値を受け取って、新しい表の値と新しい状態値を返す関数ということがわかりやすいです。
簡単な例として、

(State x) >>= f >>= g >>= h

という値を図に書いて見ます。

そのまんまです。
これでわかるように、State モナドでは、>>= (bind) でどんどんつなげていくと、横に長いラムダ(関数)を持つ State モナド値ができます。出来上がったState モナドは実質的に関数なので状態ではないですが、その関数は初期状態を待っています。モナドから関数を取り出して初期状態の値を与えると、表の値と状態値が変化していき、最終的にタプルとなって出てきます。これで大体 >>= の動作がわかりました。

最後に練習として、第6回局所的な「状態」を利用するためのStateモナド（2ページ目） | 日経 xTECH（クロステック）の例を図でかいてみます。
ここでは状態を得る get と状態を設定する put という関数が出てきます。また、>>= の代わりに >> が使われています。
get と put の定義です。

class (Monad m) => MonadState s m | m -> s where
	get :: m s
	put :: s -> m ()

instance MonadState s (State s) where
	get   = State $ \s -> (s, s)
	put s = State $ \_ -> ((), s)
第6回局所的な「状態」を利用するためのStateモナド（2ページ目） | 日経 xTECH（クロステック）

get の図はこんなふうになります。

矢印はともに恒等関数(id)です。入力がそのまま表の値と状態値に出ます。
put s は

となります。入力が途切れてますが、これは入力を捨てるということを表しています。
出力は、表の値が () 、状態が s (put の引数)です。表の値は多分、特に意味がないという意味でしょう。