State Monad and STRef

ST.scala

/**
  * Created by ariwaranosai on 16/8/28.
  *
  */

import ST._

case class World[S]()

case class ST[S, A](f: World[S] => (World[S], A)) {
  def apply(s: World[S]): (World[S], A) = f(s)

  def flatMap[B](g: A => ST[S, B]): ST[S, B] =
    ST(s => {
      val nst = f(s)
      g(nst._2).f(nst._1)
    })

  def map[B](f: A => B): ST[S, B] =
    flatMap(x => returnST(f(x)))
}

object ST {
  def empty = World[Unit]()

  def returnST[S, A](a: => A): ST[S, A] = ST(s => (s, a))

  def newVar[S, A](a: => A) = returnST[S, STRef[S, A]](STRef(a))

  type ForallST[A] = Forall[({type lambda[S] = ST[S, A]})#lambda]

  def runST[A](forall: Forall[({type lambda[S] = ST[S, A]})#lambda]): A =
    forall.apply.f(empty)._2
}


trait Forall[P[_]] {
  def apply[A]: P[A]
}

case class STRef[S, A](a: A) { self =>
  private var value: A = a

  def get: ST[S, A] = returnST(value)
  def set(a: A): ST[S, STRef[S, A]] =
    ST((s: World[S]) => {
      value = a
      (s, this)
    })

  def modify(f: A => A): ST[S, STRef[S, A]] = for {
    a <- get
    v <- set(f(a))
  } yield v
}

object Run extends App {
  def e1[S]: ST[S, STRef[S, Int]] = for {
    r <- ST.newVar[S, Int](12)
    _ <- r.modify(_ + 1)
    _ <- r.modify(_ + 1)
    _ <- r.modify(_ + 1)
  } yield r


  def e2[A] = e1[A].flatMap(_.get)

  val t = ST.runST(new ST.ForallST[Int]{def apply[A] = e2[A]})

  println(t)
}