scalaz.StateT で scalaz.State と Option を組み合わせる

scalaz の ExampleState に StateT のサンプルが無かったので試してみた。
気が向いたら、そもそも State モナドって何よ？という所から説明書くかも。

package com.github.cooldaemon.try_scalaz

import scalaz._

object TryState {
  import Scalaz._

  def run {
    def initO = stateT[Option, Int, Int]((s: Int) =>
      if (s % 2 == 0) Some((s, s)) else None
    )

    initO ! 1 assert_=== None
    initO ! 2 assert_=== Some(2)

    def putO(s: Int) = stateT[Option, Int, Unit](_ =>
      if (s % 2 == 0) Some((s, ())) else None
    )

    putO(1) ~> 0 assert_=== None
    putO(2) ~> 0 assert_=== Some(2)

    val res1 = for {
      x <- initO;
      _ <- putO(4)
    } yield x

    res1 ! 1  assert_=== None
    res1 ! 2  assert_=== Some(2)
    res1 ~> 2 assert_=== Some(4)

    def modifyO(f: Int => Int) =
      initO flatMap (s => stateT[Option, Int, Unit](_ =>
        if (s % 2 == 0) Some((f(s), ())) else None
      ))

    modifyO(s => s + 1) ~> 1 assert_=== None
    modifyO(s => s + 1) ~> 2 assert_=== Some(3)

    def getsO = stateT[Option, Int, Int]((s: Int) =>
      Some((s, s))
    )

    val res2 = for {
      _ <- modifyO(s => s + 2);
      _ <- modifyO(s => s + 1);
      x <- getsO
    } yield x

    res2 ! 1 assert_=== None
    res2 ! 2 assert_=== Some(5)

    val res3 = for {
      _ <- res1;
      x <- res2
    } yield x

    res3 ! 1 assert_=== None
    res3 ! 2 assert_=== Some(7)
  }
}