順を追って理解するStateモナドの実装

突然ですが、Stateモナドの実装を書けますか？

Stateモナドは自由に読み書き可能なグローバル変数のような機能を、副作用を伴わない純粋な関数だけを使って実装するための便利な概念です。使うだけなら簡単なので普段から便利に利用してる人も多いかもしれませんが、いざ自分で実装しろと言われると手が止まってしまう方もいるのではないでしょうか。

今回はそんなStateモナドの実装を丁寧に見ていき、仕組みを理解することを目指したいと思います。

State

さっそく以下に示すのがStateモナドの型の定義です^[1]。

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

まずStateはnewtypeを使って関数と同じ型として定義されます。この関数は以下の構造を意識して作られています。

実行前後で状態の型は変わりませんが、値に関しては変更することが可能になっているのです。

モナドにするためにはFunctorとApplicativeのインスタンスにする必要があります。

Functor

まずはFunctorのインスタンス実装を見てみましょう。

instance Functor (State s) where
  fmap f (State g) = State $ \s ->
    let (s', a) = g s
     in (s', f a)

• 2行目: 返り値であるStateの実態は関数なので値を構築する際に、"前の状態"を関数の引数として（\s ->）受け取ることができます。
• 3行目: 引数として受け取った"前の状態"を使って関数gを実行します。ただしここは関数定義の中なので、実際に実行されるのは今定義している関数が実行される時であることに注意です。gを実行した結果として次の状態s'と実行結果aを受け取ります。
• 4行目: 今構築しているStateの実行後の状態としてはs'を、そして今実装したかったのはFunctorのfmapなのでf aを実行結果としてタプルを構築しています。

関数の引数として導入した"前の状態"を使って"次の状態"を巧みに得ているのがポイントです。

Applicative

Functorと同様にApplicativeも実装することができます。

instance Applicative (State s) where
  pure a = State $ \s -> (s, a)
  (State f) <*> (State g) = State $ \s ->
    let (s', h) = f s
        (s'', a) = g s'
     in (s'', h a)

pureの型がa -> State s a, <*> の型がState s (a -> b) -> State s a -> State s bであることを思い出しておいてください。

• 2行目: pure は状態を変化することなく左から右に流しながら、与えられた値を実行結果として返します。
• 3行目: State が関数であることを利用して、引数として前の状態sを導入します
• 4行目: まずf :: s -> (s, a -> b) を"前の状態"sで評価することで、次の状態s'と実行結果h :: a -> bを得ます。
• 5行目: 4行目で得られた次の状態s'を使って更にgを評価することで、次の次の状態s''と実行結果aを得ます。
• 6行目: 次の次の状態s''と実行結果h aのタプルを作ります。

少し煩雑かもしれませんが、与えられたStateを順番に、状態を更新しながら実行していき最後に計算したかった値h aと最終的な状態s''を返しているのがポイントです。

まるで"状態"という水が流れる配管をうまくつなぎ合わせているみたいですね！

Monad

Applicativeの実装が理解できればMonadのインスタンス実装は簡単です。

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

>>= の型がState s a -> (a -> State s b) -> State s bであることを思い出しておいてください。

• 2行目: お決まりのパターンですね。State が関数であることを利用して、引数として前の状態sを導入しています。
• 3行目: 1つ目のStateの実態であるfを"前の状態"を使って評価し、次の状態s'と実行結果aを得ます。
• 4行目: モナドの場合2つ目のStateが3行目で計算した実行結果を使わないと手に入らないので、まずm aを計算しています。その結果にrunStateを行い通常の関数にしたところで、s'を使って評価し、最終的な結果を得ています。

以上でモナドの実装は終わりです。

純粋な関数だけからグローバル変数のような機能を生み出すStateモナドは魔法のようで実装は難解に思われるかもしれませんが、順を追って理解すれば"状態"の流れをうまく配線していってる様子がわかるかと思います。Stateモナドも「なんだこんなものか」と思っていただければ幸いです🤗

---

＼読んでいただきありがとうございました！／
この記事が面白かったら いいね♡ をいただけると嬉しいです☺️
バッジを贈っていただければ次の記事を書くため励みになります🙌

脚注

1. Stateモナドを含む代表的なライブラリである transformers では返り値のタプルが逆で s -> (a, s) と定義されています