state.scala

package state

case class State[S, +A](run: S => (A, S)) {
  def flatMap[AA >: A, B](f: AA => State[S, B]): State[S, B] =
    State({ s0: S =>
      val (a, s1) = run(s0)
      f(a).run(s1)
    })

  def map[AA >: A, B](f: AA => B): State[S, B] =
    flatMap(f andThen State.unit[S, B])

  def map2[AA >: A, B, C](rb: State[S, B])(f: (AA, B) => C): State[S, C] =
    flatMap((a: A) => rb.map((b: B) => f(a, b)))
}

object State {
  def unit[S, A](a: A): State[S, A] =
    State(s => (a, s))

  def sequence[S, A](rs: List[State[S, A]]): State[S, List[A]] =
    rs match {
      case Nil => unit(Nil)
      case h :: t => h.map2(sequence(t))(_ :: _)
    }

  def get[S]: State[S, S] =
    State(s => (s, s))

  def set[S](s: S): State[S, Unit] =
    State(_ => ((), s))
}