|
| 1 | +package arrow.continuations.generic |
| 2 | + |
| 3 | +import kotlinx.atomicfu.atomic |
| 4 | +import kotlinx.atomicfu.loop |
| 5 | +import kotlin.coroutines.Continuation |
| 6 | +import kotlin.coroutines.EmptyCoroutineContext |
| 7 | +import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED |
| 8 | +import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn |
| 9 | +import kotlin.coroutines.resume |
| 10 | +import kotlin.coroutines.suspendCoroutine |
| 11 | + |
| 12 | +/** |
| 13 | + * (Simulated) Multishot capable delimited control scope |
| 14 | + * |
| 15 | + * This has several drawbacks: |
| 16 | + * - f will rerun completely on multishot and only the results of [shift] are cached so any sideeffects outside of |
| 17 | + * [shift] will rerun! |
| 18 | + * - This accumulates all results of [shift] (every argument passed when invoking the continuation) so on long running computations |
| 19 | + * this may keep quite a bit of memory |
| 20 | + * - If the pure part before a multishot is expensive the multishot itself will have to rerun that, which makes it somewhat slow |
| 21 | + * - This is terribly hard to implement properly with nested scopes (which this one does not support) |
| 22 | + * |
| 23 | + * As per usual understanding of [DelimContScope] is required as I will only be commenting differences for now. |
| 24 | + */ |
| 25 | +open class MultiShotDelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedScope<R> { |
| 26 | + |
| 27 | + private val resultVar = atomic<R?>(null) |
| 28 | + private val nextShift = atomic<(suspend () -> R)?>(null) |
| 29 | + |
| 30 | + // TODO This can be append only and needs fast reversed access |
| 31 | + private val shiftFnContinuations = mutableListOf<Continuation<R>>() |
| 32 | + |
| 33 | + /** |
| 34 | + * Keep the arguments passed to [DelimitedContinuation.invoke] to be able to replay the scope if necessary |
| 35 | + */ |
| 36 | + // TODO This can be append only and needs fast random access and slicing |
| 37 | + internal open val stack = mutableListOf<Any?>() |
| 38 | + |
| 39 | + /** |
| 40 | + * Our continuation now includes the function [f] to rerun on multishot, the current live (single-shot) continuation, |
| 41 | + * the current stack and the offset from that stack when this is created which is used to know when to resume normal |
| 42 | + * execution again on a replay. |
| 43 | + */ |
| 44 | + class MultiShotCont<A, R>( |
| 45 | + liveContinuation: Continuation<A>, |
| 46 | + private val f: suspend DelimitedScope<R>.() -> R, |
| 47 | + private val stack: MutableList<Any?>, |
| 48 | + private val shiftFnContinuations: MutableList<Continuation<R>> |
| 49 | + ) : DelimitedContinuation<A, R> { |
| 50 | + // To make sure the continuation is only invoked once we put it in a nullable atomic and only access it through getAndSet |
| 51 | + private val liveContinuation = atomic<Continuation<A>?>(liveContinuation) |
| 52 | + private val stackOffset = stack.size |
| 53 | + |
| 54 | + override suspend fun invoke(a: A): R = |
| 55 | + when (val cont = liveContinuation.getAndSet(null)) { |
| 56 | + // On multishot we replay with a prefilled stack from start to the point at which this object was created |
| 57 | + // (when the shift block this runs in was first called) |
| 58 | + null -> PrefilledDelimContScope((stack.subList(0, stackOffset).toList() + a).toMutableList(), f).invoke() |
| 59 | + // on the first pass we operate like a normal delimited scope but we also save the argument to the stack before resuming |
| 60 | + else -> suspendCoroutine { resumeShift -> |
| 61 | + shiftFnContinuations.add(resumeShift) |
| 62 | + stack.add(a) |
| 63 | + cont.resume(a) |
| 64 | + } |
| 65 | + } |
| 66 | + } |
| 67 | + |
| 68 | + data class CPSCont<A, R>( |
| 69 | + private val runFunc: suspend DelimitedScope<R>.(A) -> R |
| 70 | + ) : DelimitedContinuation<A, R> { |
| 71 | + override suspend fun invoke(a: A): R = DelimContScope<R> { runFunc(a) }.invoke() |
| 72 | + } |
| 73 | + |
| 74 | + override suspend fun <A> shift(func: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A = |
| 75 | + suspendCoroutine { continueMain -> |
| 76 | + val c = MultiShotCont(continueMain, f, stack, shiftFnContinuations) |
| 77 | + assert(nextShift.compareAndSet(null, suspend { this.func(c) })) |
| 78 | + } |
| 79 | + |
| 80 | + override suspend fun <A, B> shiftCPS(func: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing = |
| 81 | + suspendCoroutine { |
| 82 | + assert(nextShift.compareAndSet(null, suspend { func(CPSCont(c)) })) |
| 83 | + } |
| 84 | + |
| 85 | + // This assumes RestrictSuspension or at least assumes the user to never reference the parent scope in f. |
| 86 | + override suspend fun <A> reset(f: suspend DelimitedScope<A>.() -> A): A = |
| 87 | + MultiShotDelimContScope(f).invoke() |
| 88 | + |
| 89 | + fun invoke(): R { |
| 90 | + f.startCoroutineUninterceptedOrReturn(this, Continuation(EmptyCoroutineContext) { result -> |
| 91 | + resultVar.value = result.getOrThrow() |
| 92 | + }).let { |
| 93 | + if (it == COROUTINE_SUSPENDED) { |
| 94 | + resultVar.loop { mRes -> |
| 95 | + if (mRes == null) { |
| 96 | + val nextShiftFn = nextShift.getAndSet(null) |
| 97 | + ?: throw IllegalStateException("No further work to do but also no result!") |
| 98 | + nextShiftFn.startCoroutineUninterceptedOrReturn(Continuation(EmptyCoroutineContext) { result -> |
| 99 | + resultVar.value = result.getOrThrow() |
| 100 | + }).let { |
| 101 | + if (it != COROUTINE_SUSPENDED) resultVar.value = it as R |
| 102 | + } |
| 103 | + } else return@let |
| 104 | + } |
| 105 | + } else return@invoke it as R |
| 106 | + } |
| 107 | + assert(resultVar.value != null) |
| 108 | + for (c in shiftFnContinuations.asReversed()) c.resume(resultVar.value!!) |
| 109 | + return resultVar.value!! |
| 110 | + } |
| 111 | + |
| 112 | + companion object { |
| 113 | + fun <R> reset(f: suspend DelimitedScope<R>.() -> R): R = MultiShotDelimContScope(f).invoke() |
| 114 | + } |
| 115 | +} |
| 116 | + |
| 117 | +class PrefilledDelimContScope<R>( |
| 118 | + override val stack: MutableList<Any?>, |
| 119 | + f: suspend DelimitedScope<R>.() -> R |
| 120 | +) : MultiShotDelimContScope<R>(f) { |
| 121 | + var depth = 0 |
| 122 | + |
| 123 | + // Here we first check if we still have values in our local stack and if so we use those first |
| 124 | + // if not we delegate to the normal delimited control implementation |
| 125 | + override suspend fun <A> shift(func: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A = |
| 126 | + if (stack.size > depth) stack[depth++] as A |
| 127 | + else super.shift(func).also { depth++ } |
| 128 | +} |
0 commit comments