local_state.scala

// This new data type will employ Scala’s type system to gain two static guarantees. We want our code to not compile if it violates these invariants:
// - If we hold a reference to a mutable object, then nothing can observe us mutat- ing it.
// - A mutable object can never be observed outside of the scope in which it was created.

/**
 * For state thread, state transi- tion, state token, or state tag
 * A local-effects monad.
 *
 * A: type of the mutable value.
 * S: represents the ability to mutate state
 */
trait ST[S,A] { self =>
  protected def run(s: S): (A,S)

  def map[B](f: A => B): ST[S,B] = new ST[S,B] {
    def run(s: S) = {
      val (a, s1) = self.run(s)
      (f(a), s1)
    }
  }

  def flatMap[B](f: A => ST[S,B]): ST[S,B] = new ST[S,B] {
    def run(s: S) = {
      val (a, s1) = self.run(s)
      f(a).run(s1)
    }
  }

}

object ST {
  def apply[S,A](a: => A) = {
    lazy val memo = a
    new ST[S,A] { def run(s: S) = (memo, s) }
  }

  def runST[A](st: RunnableST[A]): A =
    st.apply[Unit].run(())._1

}

/**
 * A: the type of the mutable value in STRef.
 * S: never used, but needed to create a ST type
 */
sealed trait STRef[S, A] {
  protected var cell: A

  def read: ST[S, A] = ST(cell)

  def write(a: A): ST[S, Unit] = new ST[S,Unit] {

  def run(s: S) = { cell = a
    ((), s)
  }}
}

object STRef {
  def apply[S,A](a: A): ST[S, STRef[S, A]] = ST(new STRef[S, A] {
    var cell = a
  })
}

/**
 * A trait to prevent forbidden states :
 * - An ST action cannot return a mutable state (a STRef)
 * - any ST[S,T] where T involves the type S.
 */
trait RunnableST[A] {
  def apply[S]: ST[S,A]
}

sealed abstract class STArray[S,A](implicit manifest: Manifest[A]) {
  protected def value: Array[A]

  def size: ST[S,Int] = ST(value.size)

  def write(i: Int, a: A): ST[S,Unit] = new ST[S,Unit] { def run(s: S) = {
    value(i) = a
    ((), s) }
  }

  def read(i: Int): ST[S,A] = ST(value(i))

  def freeze: ST[S,List[A]] = ST(value.toList)

  def fill(xs: Map[Int,A]): ST[S,Unit] = {
    xs.foldRight(ST[S,Unit](())) {
      case ((k, v), st) => st flatMap (_ => write(k, v))
    }
  }
}

object STArray {
  def apply[S,A:Manifest](sz: Int, v: A): ST[S, STArray[S,A]] =
    ST(new STArray[S,A] {
      lazy val value = Array.fill(sz)(v)
    })}