Delimited Continuations shift/reset in Kotlin

Delimited.kt

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

/**
 * Implementation for Delimited Continuations `shift`/`reset` primitives via Kotlin Coroutines.
 * See [https://en.wikipedia.org/wiki/Delimited_continuation].
 *
 * The following LISP code:
 *
 * ```
 * (* 2 (reset (+ 1 (shift k (k 5)))))
 * ```
 *
 * translates to:
 *
 * ```
 * 2 * reset<Int> {
 *     1 + shift<Int> { k -> k(5) }
 * }
 * ```
 */
fun <T> reset(body: suspend DelimitedScope<T>.() -> T): T =
    DelimitedScopeImpl<T>().also { impl ->
        body.startCoroutine(impl, impl)
    }.runReset()

interface DelimitedContinuation<T, R>

@RestrictsSuspension
abstract class DelimitedScope<T> {
    abstract suspend fun <R> shift(block: suspend DelimitedScope<T>.(DelimitedContinuation<T, R>) -> T): R
    abstract suspend operator fun <R> DelimitedContinuation<T, R>.invoke(value: R): T
}

private typealias ShiftedFun<T> = (DelimitedScope<T>, DelimitedContinuation<T, Any?>, Continuation<T>) -> Any?

@Suppress("UNCHECKED_CAST")
private class DelimitedScopeImpl<T> : DelimitedScope<T>(), Continuation<T>, DelimitedContinuation<T, Any?> {
    private var shifted: ShiftedFun<T>? = null
    private var shiftCont: Continuation<Any?>? = null
    private var invokeCont: Continuation<T>? = null
    private var invokeValue: Any? = null
    private var result: Result<T>? = null

    override val context: CoroutineContext
        get() = EmptyCoroutineContext

    override fun resumeWith(result: Result<T>) {
        this.result = result
    }

    override suspend fun <R> shift(block: suspend DelimitedScope<T>.(DelimitedContinuation<T, R>) -> T): R =
        suspendCoroutineUninterceptedOrReturn {
            this.shifted = block as ShiftedFun<T>
            this.shiftCont = it as Continuation<Any?>
            COROUTINE_SUSPENDED
        }

    override suspend fun <R> DelimitedContinuation<T, R>.invoke(value: R): T =
        suspendCoroutineUninterceptedOrReturn sc@{
            check(invokeCont == null)
            invokeCont = it
            invokeValue = value
            COROUTINE_SUSPENDED
        }

    fun runReset(): T {
        // This is the stack of continuation in the `shift { ... }` after call to delimited continuation
        var currentCont: Continuation<T> = this
        // Trampoline loop to avoid call stack usage
        loop@while (true) {
            // Call shift { ... } body or break if there are no more shift calls
            val shifted = takeShifted() ?: break
            // If shift does not call any continuation, then its value becomes the result -- break out of the loop
            try {
                val value = shifted.invoke(this, this, currentCont)
                if (value !== COROUTINE_SUSPENDED) {
                    result = Result.success(value as T)
                    break
                }
            } catch (e: Throwable) {
                result = Result.failure(e)
                break
            }
            // Shift has suspended - check if shift { ... } body had invoked continuation
            currentCont = takeInvokeCont() ?: continue@loop
            val shiftCont = takeShiftCont()
                ?: error("Delimited continuation is single-shot and cannot be invoked twice")
            shiftCont.resume(invokeValue)
        }
        // Propagate the result to all pending continuations in shift { ... } bodies
        currentCont.resumeWith(result!!)
        // Return the final result
        return result!!.getOrThrow()
    }

    private fun takeShifted() = shifted?.also { shifted = null }
    private fun takeShiftCont() = shiftCont?.also { shiftCont = null }
    private fun takeInvokeCont() = invokeCont?.also { invokeCont = null  }
}

DelimitedTest.kt

import org.junit.*
import kotlin.test.*

class DelimitedTest {
    @Test
    fun testNoShit() {
        val x = reset<Int> { 42 }
        assertEquals(42, x)
    }

    @Test
    fun testShiftOnly() {
        val x = reset<Int> {
            shift<Int> { k -> k(42) }
        }
        assertEquals(42, x)
    }

    @Test
    fun testShiftRight() {
        val x = reset<Int> {
            40 + shift<Int> { k -> k(2) }
        }
        assertEquals(42, x)
    }

    @Test
    fun testShiftLeft() {
        val x = reset<Int> {
            shift<Int> { k -> k(40) } + 2
        }
        assertEquals(42, x)
    }

    @Test
    fun testShiftBoth() {
        val x = reset<Int> {
            shift<Int> { k -> k(40) } +
                    shift<Int> { k -> k(2) }
        }
        assertEquals(42, x)
    }

    @Test
    fun testShiftToString() {
        val x = reset<String> {
            shift<Int> { k -> k(42) }.toString()
        }
        assertEquals("42", x)
    }

    @Test
    fun testShiftWithoutContinuationInvoke() {
        val x = reset<Int> {
            shift<String> {
                42 // does not call continuation, just override result
            }
            0 // this is not called
        }
        assertEquals(42, x)
    }

    // From: https://en.wikipedia.org/wiki/Delimited_continuation
    // (* 2 (reset (+ 1 (shift k (k 5)))))
    // k := (+ 1 [])
    @Test
    fun testWikiSample() {
        val x = 2 * reset<Int> {
            1 + shift<Int> { k -> k(5) }
        }
        assertEquals(12, x)
    }

    // It must be extension on DelimitedScope<Int> to be able to shift
    private suspend fun DelimitedScope<Int>.shiftFun(x: Int): Int =
        shift<Int> { k -> k(x) } * 2

    @Test
    fun testShiftFromFunction() {
        val x = reset<Int> {
            2 + shiftFun(20)
        }
        assertEquals(42, x)
    }

    @Test
    // Ensure there's no stack overflow because of many "shift" calls
    fun testManyShifts() {
        val res = reset<String> {
            for (x in 0..10000) {
                shift<Int> { k ->
                    k(x)
                }
            }
            "OK"
        }
        assertEquals("OK", res)
    }

    @Test
    // See https://gist.github.com/elizarov/5bbbe5a3b88985ae577d8ec3706e85ef#gistcomment-3304535
    fun testShiftRemainderCalled() {
        val log = ArrayList<String>()
        val x = reset<Int> {
            val y = shift<Int> { k ->
                log += "before 1"
                val r = k(1)
                log += "after 1"
                r
            }
            log += y.toString()
            val z = shift<Int> { k ->
                log += "before 2"
                val r = k(2)
                log += "after 2"
                r
            }
            log += z.toString()
            y + z
        }
        assertEquals(3, x)
        assertEquals(listOf(
            "before 1",
            "1",
            "before 2",
            "2",
            "after 2",
            "after 1"
        ), log)
    }
}

@raulraja this is not addressing your problem at hand, but maybe you are also interested in my old monadic reflection experiments

RE your problem: To me storing the continuation in a mutable field is confusing since in the case of multiple resumptions you are running into aliasing problems.

Personally, I am a big fan of continuations as immutable liked lists of immutable frames
since it makes sharing a lot easier.

@b-studios thanks for the info. would you mind posting the link again to 'monadic reflection experiments'? it's linking to my comment on the gist. Thanks!

@raulraja sorry for posting the wrong link, I updated it

thanks @b-studios!

@b-studios looks great and I'm gonna give it a shot to see if I can make it work with coroutines and suspension. This uses loom continuations but we can't depend on that.

In our current case what here is modeled as a program being a Prompt and its body, in our case we need to deal with suspension.
You may be already aware but even if we cast a suspend function to a regular function is still goes throw the internals of the compiler-generated state machine checking for completion. https://github.com/b-studios/cats-reflect/blob/master/src/main/scala/cats/reflect/internal/Coroutine.scala#L9 I'm gonna give it a shot and see if I can model that with suspending functions.

Thanks so much for all your help and impressive research in algebraic effects and gists! Had fun and learned a lot from them.

@b-studios I've fixed implementation so that it is correct and does not use the call-stack. All the calls are now in a single top-level runReset loop. All the inner shift and continuation calls just suspend and relinquish control to this top-level loop. This way you can shift many times without using call-stack and if you have any code in shift { .... } body after continuation invocation, then the corresponding "after invocation continuations" are stored in the heap and are called in the correct order at the end.

To be clear, this implementation is still not a 100% port of shift/reset from LISP: you cannot call a continuation an arbitrary number of times, nor you can save a continuation somewhere and call it at an arbitrary moment of time later. It only supports calling a continuation from inside of shift at most once.

you cannot call a continuation an arbitrary number of times

I was about to point that out :) But still it is very cool to see that you managed to implement it in a stack-safe manner!

BTW in our design of the Effekt language we also treat resume (that is, the continuation) as an "effect" (or block, in our lingo). If you are interested, you can read up on it in our paper

@b-studios Thanks for the paper. If you squint well enough, you might notice similarities with context-based programming practice that is growing in Kotlin ecosystem due to the Kotlin's convenient syntax for functions with a receiver (context). Kotlin coroutines (suspending functions) are also context-based, as well as composable functions form JetPack Compose. Another Kotlin-specific enabling feature for this programming style are inline functions, which serve as "co-pure" functions. Yes, they are second-class, but in return they automatically propagate caller's context without having to use effect polymorphism in their signature (indeed, effect polymorphism is no-go for a general-purpose non-research language).

For pragmatic reasons, though, Kotlin approach does not readily generalize to arbitrary effects, since we a looking for an efficient runtime mapping. All the various complex contexts (like suspending context or composable context) cannot be simply implemented as a library in Kotlin, but need context-specific support from the compiler. It also drives various pragmatic restrictions like single-shot-only continuations (creating an immutable snapshot of continuation at each suspending point would be simply too expensive for many real-life use-cases we were designing Kotlin coroutines for).

@b-studios @elizarov thanks for the conversation and proposed suggestions. Jannis, Simon, and I worked in attempting multi-prompt and similar techniques and we have now a version of reset / shift that is gonna allow us to build effect handlers and remove the state label reflection tricks in the Arrow continuations. In case you are interested, this will keep evolving but this is the first attempt https://github.com/arrow-kt/arrow-core/pull/226/files if you have any comments on it feel free to do directly in the PR even if closed and we will address them there or as issues. Thanks again for all your help.

@elizarov I figured testNoShit is named incorrectly

@raulraja I'm curious how the implementation of delimited multi-shot continuation has evolved in Arrow. I traced your PR in arrow-core to a squashed commit in arrow, but I couldn't find the code on main.

@steshaw I believe that code ended up getting removed through the many evolution steps of Arrow. IMO, some of it was "too hacky" which resulted in it getting removed.
A spiritual successor to that code (although it shares almost none of its implementation) is my so-far-successful attempt at building multiprompt delimited continuations in Kotlin. It theoretically has a single-shot core that uses no "magic" (i.e reflection or platform-specific code) but I haven't extracted that yet. It uses a tiny API surface to shallow clone continuations for multishot cases (that surface just being val Continuation<*>.completion: Continuation<*> and fun <T> Continuation<T>.copy(completion: Continuation<*>): Continuation<T>, but the former isn't even quite "necessary" if we change some of the behaviour of the latter, but it's useful performance-wise). It currently works only for JVM and JS because I can't figure out how to shallow-clone on other platforms yet.
You can find some of the original Arrow test code in my repo's tests.
I'm pretty convinced that my implementation is nearly-correct. It passes a lot of non-obvious tests and seems to demonstrate the multiprompt behaviour, and hence any bugs are likely to be squashable easily. The implementation is based on @b-studios 's java-effekt but with some Kotlin-specific nuances, and a bunch of optimizations for single-shot resumptions. In fact, the repo has further examples of adding a thin layer on multiprompt delimited continuations to turn them into "Lexical Effect Handlers", which is the idea that java-effekt presents. Further, the language Effekt evolved directly from that, and pretty much every one of its examples can thus be translated directly to Kotlin (check out the effekt/casestudies tests folder for that)
I'd be glad to hear your feedback if you choose to play around with it :)