Marco's question in slack

st.kt

package arrow.prelude.test

import arrow.continuations.generic.DelimContScope
import arrow.continuations.generic.DelimitedScope

abstract class ST<S, A> @PublishedApi internal constructor() {
  companion object {
    operator fun <S, A> invoke(a: () -> A): ST<S, A> {
      val memo by lazy(a)
      return object : ST<S, A>() {
        override suspend fun run(s: S) = Pair(memo, s)
      }
    }

    suspend fun <A> runST(st: RunnableST<A>): A =
      st.invoke<Unit>().run(Unit).first
  }

  @PublishedApi
  internal abstract suspend fun run(s: S): Pair<A, S>

  inline fun <B> map(crossinline f: (A) -> B): ST<S, B> = object : ST<S, B>() {
    override suspend fun run(s: S): Pair<B, S> {
      val (a, s1) = this@ST.run(s)
      return Pair(f(a), s1)
    }
  }

  fun <B> flatMap(f: suspend (A) -> ST<S, B>): ST<S, B> = object : ST<S, B>() {
    override suspend fun run(s: S): Pair<B, S> {
      val (a, s1) = this@ST.run(s)
      return f(a).run(s1)
    }
  }
}

interface RunnableST<A> {
  suspend fun <S> invoke(): ST<S, A>
}

abstract class STRef<S, A> private constructor() {
  companion object {
    operator fun <S, A> invoke(a: A): ST<S, STRef<S, A>> = ST {
      object : STRef<S, A>() {
        override var cell: A = a
      }
    }
  }

  protected abstract var cell: A

  fun read(): ST<S, A> = ST {
    cell
  }

  fun write(a: A): ST<S, Unit> = object : ST<S, Unit>() {
    override suspend fun run(s: S): Pair<Unit, S> {
      cell = a //I guess this is atomic?
      return Pair(Unit, s)
    }
  }
}

/**
 * Effect will replace [Monad] an all functor like combinators that need
 * kind emulation since in suspension all those are unnecessary given you
 * can always obtain F<A> -> A
 */
interface Effect<F> {
  val scope: DelimitedScope<F>
}

/**
 * A handler implements a set of capabilities for a use case.
 * In this case [EitherEffect] serves the purpose of describing what's available
 * in [st] computational blocks. This declaration is equivalent to a
 * type class instances if we accept the block is how type-classes expose all capabilities.
 * While these blocks are monadic in this implementation all kinds of other specialized
 * handlers are possible.
 */
class STEffect<S, A>(override val scope: DelimitedScope<ST<S, A>>) :
  Effect<ST<S, A>> {

  suspend operator fun <B> ST<S, B>.invoke(): B =
    scope.shift { cb -> flatMap { cb(it) } }
}

/**
 * This function exposes your effect for users to consume where they can access all functions
 * described in in [STEffect]
 */
inline fun <S, A> st(crossinline f: suspend STEffect<S, A>.() -> A): ST<S, A> =
  DelimContScope.reset { f(STEffect(this)).run { ST { this } } }

//TODO: wish to write the following with a for comprehension
suspend fun x() = STRef<String, Int>(10).flatMap { r1 ->
  STRef<String, Int>(20).flatMap { r2 ->
    r1.read().flatMap { x ->
      r2.read().flatMap { y ->
        r1.write(y + 1).flatMap {
          r2.write(x + 1).flatMap {
            r1.read().flatMap { a ->
              r2.read().map { b ->
                Pair(a, b)
              }
            }
          }
        }
      }
    }
  }
}

suspend fun xSimpler(): ST<String, Pair<Int, Int>> =
  st {
    val r1 = STRef<String, Int>(10)()
    val r2 = STRef<String, Int>(10)()
    val x = r1.read()()
    val y = r2.read()()
    r1.write(y + 1)()
    r2.write(x + 1)()
    val a = r1.read()()
    val b = r2.read()()
    a to b
}

suspend fun main() {
  println(x().run("flatMap"))
  println(xSimpler().run("comprehension"))
}