ST Monad in Kotlin, with subregioning, and with si...

# functional

[Youssef Shoaib [MOD]]

ST Monad in Kotlin, with subregioning, and with similar safety guarantees as in Haskell! (STRef and the STScope itself may both escape, but they're only usable exactly in the region they were gotten in!) (playground)

import kotlin.coroutines.Continuation
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.coroutines.RestrictsSuspension
import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn

class STRef<in R, S> internal constructor(internal var state: S)

@RestrictsSuspension
interface STScope<R> {
  suspend fun <S> new(init: S): STRef<R, S> = STRef(init)
  suspend fun <S> STRef<R, S>.get(): S = state
  suspend fun <S> STRef<R, S>.set(value: S) {
    state = value
  }
  suspend fun subregion(block: suspend SubSTScope<*, R>.() -> Unit)
}

// could use typealias instead, but bad IDE support
interface SubSTScope<R1 : R2, R2> : STScope<R1>

class STScopeImpl<R> : SubSTScope<R, R> {
  // could even be inline!
  override suspend fun subregion(block: suspend SubSTScope<*, R>.() -> Unit) = block()
}

fun <R> runST(block: suspend STScope<*>.() -> R): R =
  block.startCoroutineUninterceptedOrReturn(STScopeImpl<Any?>(), Continuation(EmptyCoroutineContext) {}) as R

fun main() {
  runST {
    // need to use inner functions here because the compiler isn't smart enough about existential
    // Outference might solve this problem!
    suspend fun <R> STScope<R>.example() {
      val ref = new(0)
      ref.set(42)
      println(ref.get()) // prints 42
      //val escape = { ref.get() } // not allowed to escape
      //val escape = runST { ref.get() } // still not allowed to escape
      //runST { with(this@example) { ref.get() } } // even this doesn't work!
      subregion {
        suspend fun <R1: R> STScope<R1>.inside() {
          val innerRef = new(0)
          innerRef.set(100)
          println(innerRef.get()) // prints 100
          println(ref.get()) // prints 42
        }
        inside()
      }
    }
    example()
  }
}

I hence conjecture that @RestrictsSuspension subsumes the concept of monadic regions, albeit while limiting some of Kotlin's magic (can't use general suspend functions inside, can't extend them easily, etc).
In fact, I think Kotlin may have just enough features to enable general capability/resource tracking.

[Youssef Shoaib [MOD]]

A simple File example, with again the same safety guarantees (file handles cannot be used outside the region they're created in, hence you can never read from a closed file!):

import Example.example
import java.nio.file.Path
import kotlin.coroutines.Continuation
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.coroutines.RestrictsSuspension
import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn
import kotlin.io.path.reader

private object Example {
  class FileHandle<R> internal constructor(internal val reader: java.io.Reader)

  interface AutoCloseScope<R> {
    fun onClose(action: () -> Unit)
  }

  @RestrictsSuspension
  sealed interface Region<R> : AutoCloseScope<R>

  private class RegionImpl<R> : Region<R>, AutoCloseable {
    private val closeActions = mutableListOf<() -> Unit>()
    override fun onClose(action: () -> Unit) {
      closeActions.add(action)
    }

    override fun close() {
      for (action in closeActions) {
        try {
          action()
        } catch (e: Exception) {
          // Handle exceptions from close actions if necessary
          e.printStackTrace()
        }
      }
      closeActions.clear()
    }
  }

  fun <R> region(block: suspend Region<*>.() -> R): R = RegionImpl<Any?>().use {
    block.startCoroutineUninterceptedOrReturn(it, Continuation(EmptyCoroutineContext) {}) as R
  }
  suspend fun <R, T> Region<R>.subregion(block: suspend Region<out R>.() -> T): T = RegionImpl<R>().use { block(it) }

  context(scope: AutoCloseScope<R>)
  suspend fun <R> Region<out R>.open(path: Path): FileHandle<R> =
    FileHandle(path.reader().also { scope.onClose(it::close) })

  suspend fun <R> Region<out R>.read(file: FileHandle<R>): Char = file.reader.read().toChar()

  suspend fun <R> Region<R>.example() {
    val file = open(Path.of("example.txt"))
    println(read(file))
    val file2 = subregion {
      suspend fun <R1 : R> Region<R1>.subExample(): FileHandle<R> {
        val subfile = open(Path.of("example2.txt"))
        println(read(file))
        println(read(subfile))
        return open<R>(Path.of("example3.txt"))
      }
      subExample()
    }
    read(file2)
  }
}

private fun main() = Example.region {
  example()
}

Edit: I cleaned up the file example, and made it a general auto-closing region. I also fixed some edge cases w.r.t. being able to open a file that's managed by a parent region even if you're in a subregion