Stateモナドを理解する

(元記事は http://ku00.hatenablog.com/entry/2017/06/04/165151 です)

<blockquote class="twitter-tweet" data-lang="ja"><p lang="ja" dir="ltr">Stateモナド、定期的にわかった→わからんを繰り返してる</p>&mdash; わさびのり太郎 (@ku00_) <a href="https://twitter.com/ku00_/status/869141588252909568">2017年5月29日</a></blockquote>
<script async src="//platform.twitter.com/widgets.js" charset="utf-8"></script>

ということがあったので主に備忘録としてまとめました。

この記事では [『すごいHaskellたのしく学ぼう』](http://amzn.asia/1Ol5IzO) (通称「すごいH本」)からコイントスをシミュレーションする関数を書く問題を参考に、Stateモナドについて学んでいこうと思います。


## はじめに

下記のような3回コインを投げてその裏表を返す関数があったとします。

```haskell
threeCoins :: StdGen -> (Bool, Bool, Bool)
threeCoins gen =
    let (c1, g1) = random gen :: (Bool, StdGen)
        (c2, g2) = random g1 :: (Bool, StdGen)
        (c3, g3) = random g2 :: (Bool, StdGen)
    in  (c1, c2, c3)
```

ところがStateモナドを使うことで乱数の状態をより簡単に扱うことができるようになります。

```haskell
randomSt :: (RandomGen g, Random a) => State g a
randomSt = state random

threeCoins' :: State StdGen (Bool, Bool, Bool)
threeCoins' = do
  c1 <- randomSt
  c2 <- randomSt
  c3 <- randomSt
  return (c1, c2, c3)
```

Stateモナドとは一体どのようなものなのでしょうか。またそれは通常の関数による値の受け渡しとどう違い、どのように便利なのか。それを知るためにはまず状態を伴う関数について考える必要があります。


## 状態を伴う関数

System.Randomのrandom関数は、乱数ジェネレータを引数に取り、乱数と新しいジェネレータを返します。random関数が乱数ジェネレータを受け取って新しいジェネレータを返しているのは、ジェネレータの状態を上書きすることができないからです。これは純粋関数型言語の参照透過性によるものです。

```haskell
random :: (RandomGen g, Random a) => g -> (a, g)
```

そんなrandom関数の型から、 **状態付き計算** という汎用的な型を得ることができます。
状態付き計算は、状態をsとし、計算結果aと更新された状態を返す関数として下記のように表現できます。

```
s -> (a, s)
```

## Stateモナド

状態付き計算をモナドとして扱えるのがStateモナドです。
Stateモナドは次のような定義となっています。

```haskell
type State s a = s -> (a, s)
```

Stateモナドは値コンストラクタ( `State` )を提供していないので、Stateモナドで状態付き計算を包みたい場合はstate関数を使う必要があります。

### Stateモナドの関数たち

提供されている関数は以下のようなものがあります。

- state
- runState
- evalState
- execState
- get
- put

state関数は状態付き計算を受け取ってStateモナドで包んで返します。

```
> :t state
state :: (s -> (a, s)) -> State s a
> :t state $ \x -> (25, x)
state $ \x -> (25, x) :: (Num a) => State s a
```

runState関数はStateモナドと状態を受け取って、Stateモナドを解いて状態付き計算を返します。state関数の逆ですね。

```
> :t runState
runState :: State s a -> s -> (a, s)
> runState (state $ \x -> (25, x)) "neko"
(25,"neko")
```

evalState/execState関数は計算結果か状態のどちらかだけを返したいときに使います。

```
> :t evalState
evalState :: State s a -> s -> a
> :t execState
execState :: State s a -> s -> s
```

get・put関数は、コイントス問題とは直接関係ないので後で説明します。


### returnと>>=

returnと>>=の実装があるStateモナドのインスタンス定義は次のようになっています。

```haskell
instance Monad (State s) where
  return x = State $ \s -> (x, s)
  (State h) >>= f = State $ \s -> let (a, newState) = h s
                                      (State g) = f a
                                  in g newState
```

returnは値を受け取ってそれが結果になるような状態付き計算を返すだけです。

```haskell
> :t (return "neko") :: State s String
(return "neko") :: State s String :: State s String
```

Stateモナドを>>=で関数を適用した場合はその結果もまたStateモナドである必要があるため、右辺の始まりは `State $ \s -> ...` となっています。
まずは **状態付き計算** である左辺のhと、右辺のラムダ式で引数として渡されている **状態** sを使って計算結果aと新しい状態newStateのペアを返しています。
次に左辺の関数fを計算結果aに適用すると新しい **状態付き計算** gを返します。
あとは **状態付き計算** gにnewStateを適用すれば最終的な計算結果と状態のペアを得ることができます。

これでコイントス問題をStateモナドで解くのに必要なことは説明したのでもう一度考えてみましょう。


## Stateモナドを使ってコイントス問題を解く

コイントス問題に使われるrandom関数は乱数ジェネレータを受け取って乱数と新しいジェネレータを返すのでした。これをStateモナドとして扱うためにはstate関数を使って包めばよさそうです。

```haskell
randomSt :: (RandomGen g, Random a) => State g a
randomSt = state random
```

あとはrandomSt関数を呼んでいって結果を受け取り最後にreturnで返すだけです。

```haskell
threeCoins' :: State StdGen (Bool, Bool, Bool)
threeCoins' = do
  c1 <- randomSt
  c2 <- randomSt
  c3 <- randomSt
  return (c1, c2, c3)
```

なぜrandomSt関数に状態付き計算を明示的に渡す必要がないのかは、ラムダ式を使って書いてみると理解できると思います。

```haskell
threeCoins' = randomSt >>= (\c1 -> randomSt c1 >>= (\c2 -> randomSt c2 >>= return (c1, c2, c3)))
```

runState関数などに適用して実際に使ってみましょう。

```
> runState threeCoins' $ mkStdGen 7
((True,False,True),33684305 2103410263)
```

本編はこれで終わりです。あとは上で説明しなかったStateモナドの関数の説明だけします。


## get・put関数

コイントス問題では使いませんでしたが、get・put関数も重要な関数です。

get関数は現在の状態をそのまま計算結果としたStateモナドを返します。

```
> runState get [1..5]
([1,2,3,4,5],[1,2,3,4,5])
```

put関数は状態を受け取って、その値で状態を更新したStateモナドを返します。

```
> runState (do x <- get; put $ (*2) <$> x) [1..5]
((),[2,4,6,8,10])
```


## まとめ

StateモナドはMaybeやリストモナドと比べると特殊なモナドに見えるので敬遠していましたが、Stateモナド無しで書いた関数をStateモナドを使って書き直してみるとその良さが分かるかと思います。自分は特にdo記法によるStateモナドの受け渡しをちゃんと理解していなくて、今回>>=で置き換えることをやったおかげでようやく理解ができた感じでした。

あとはStateモナドの定義が実は `type State s a = s -> (a, s)` じゃなくて `StateT s Identity` だとか、state関数はMonadState型クラスの制約があるだとか色々と省いている部分がありますがそれはまた別の機会にまとめます。今回は分かりやすさを重視しました。ソースコードなど詳しくは [Control.Monad.State.Strict](https://hackage.haskell.org/package/mtl-2.2.1/docs/Control-Monad-State-Strict.html) あたりに書いてあります。