state monad in scala?????

state.scala

case class State[S, A](runState: S => (A, S)) {
  def flatMap[B](f: A => State[S, B]): State[S, B] = State({
    s: S =>
      val (res, s1) = runState(s)
      f(res).runState(s1)
  })

  def >>=[B](f: A => State[S, B]): State[S, B] = flatMap(f)
  def >>[B](f: State[S, B]): State[S, B] = flatMap({ _ => f })

  def execState(s: S): A = {
    val (a, _) = runState(s)
    a
  }
}

object State {
  def pure[S](init: S) = State({s: S => (init, init)})

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

  def modify[S](f: S => S): State[S, Unit] = State({
    s: S => ((), f(s))
  })

  def put[S](s1: S): State[S, Unit] = modify({ _ => s1 })
}