Recursive computation engine for Kotlin

Rec.kt

import kotlin.coroutines.*
import kotlin.coroutines.intrinsics.*

/**
 * Recursive computation engine for Kotlin that uses suspending functions to keep computation
 * stack on the heap (as opposed to the real stack).
 *
 * This implementation is hardcoded for functions with two parameters [V1] and [V2] that
 * return [R] as result, but it can be easily modified for functions with different
 * number of parameters (see code).
 *
 * For example, to do DFS to find depth of a graph that is encoded by
 * `head`, `to`, and `next` arrays write:
 *
 * ```
 * // define recursive function 
 * val dfs = Rec<Int, Int, Int> { v, p ->
 *   var depth = 0
 *   var e = head[v]
 *   while (e != -1) {
 *      // use rec(...) to call this function recursively
 *      if (to[e] != p) depth = maxOf(depth, rec(to[e], v))
 *      e = next[e]
 *   }
 *   depth + 1
 * }
 *
 * // call recursive function
 * f(0, -1)
 * ```
 *
 * Note: this simplified implementation does not properly support exceptions.
 */
@Suppress("UNCHECKED_CAST")
class Rec<V1, V2, R>(
    private val block: suspend RecScope<V1, V2, R>.(V1, V2) -> R
) : RecScope<V1, V2, R>(), Continuation<R> {
    private var result: Any? = null
    private var cont: Continuation<R>? = null
    private var v1: V1? = null
    private var v2: V2? = null

    override val context: CoroutineContext
        get() = EmptyCoroutineContext

    override fun resumeWith(result: Result<R>) {
        this.result = result.getOrThrow()
        this.cont = null
    }

    operator fun invoke(v1: V1, v2: V2): R {
        set(this, v1, v2)
        run()
        return result as R
    }

    private fun run() {
        val f = block as Function4<Any?, V1?, V2?, Continuation<R>?, Any?>
        while (true) {
            val cont = this.cont ?: break
            val r = f(this, v1, v2, cont)
            if (r !== COROUTINE_SUSPENDED) cont.resume(r as R)
        }
    }

    override suspend fun rec(v1: V1, v2: V2): R = suspendCoroutineUninterceptedOrReturn { cont ->
        set(cont, v1, v2)
        COROUTINE_SUSPENDED
    }

    private fun set(cont: Continuation<R>, v1: V1, v2: V2) {
        this.cont = cont
        this.v1 = v1
        this.v2 = v2
    }
}

@RestrictsSuspension
abstract class RecScope<V1, V2, R> {
    abstract suspend fun rec(v1: V1, v2: V2): R
}