State Thread

ST.scala

opaque type ST[S, A] = S => (A, S)
object ST {
  extension [S, A](self: ST[S, A]) {
    def map[B](f: A => B): ST[S, B] = s => {
      val (a, s1) = self(s)
      (f(a), s1)
    }

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

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

  def run[A](st: [s] => () => ST[s, A]): A = {
    /* 
     * String or Int can be used in place of Unit. However, it might cause confusion as there are many values for
     * String and Int e.g.
     *   val su: ST[String, A] = st()   
     *   su("abc")._1
     */
    val su: ST[Unit, A] = st()   
    su(())._1
  }

  def lift[S, A](f: S => (A, S)): ST[S, A] = f
}

final class STRef[S, A] private (private var cell: A) {
  def read: ST[S, A]              = ST(cell)
  def write(a: => A): ST[S, Unit] = ST.lift[S, Unit] { s =>
    cell = a
    ((), s)
  }
}

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

object Main {
  @main
  def runSTApp(): Unit = {
    val p: [s] => () => ST[s, (Int, Int)] = [s] =>
      () =>
        for {
          r1 <- STRef[s, Int](1)
          r2 <- STRef[s, Int](2)
          x <- r1.read
          y <- r2.read
          _ <- r1.write(y + 1)
          _ <- r2.write(x + 1)
          a <- r1.read
          b <- r2.read
        }
        yield (a, b)

    val q: [s] => () => ST[s, STRef[s, Int]] = [s] =>
      () => STRef(1)

    val r = ST.run(p)
    // not possible to fit any type for the S position in STRef[S, Int]
    // val r2 = ST.run[STRef[Nothing, Int]](q)
    println(r)
  }
}