Delimited continuations (#226)

Co-authored-by: Raul Raja <[email protected]>
Co-authored-by: Simon Vergauwen <[email protected]>
Co-authored-by: Rachel M. Carmena <[email protected]>

Author: 4 people

arrow-libs/core/arrow-continuations/build.gradle

@@ -0,0 +1,16 @@
+plugins {
+    id "org.jetbrains.kotlin.jvm"
+    id "org.jetbrains.kotlin.kapt"
+    id "org.jlleitschuh.gradle.ktlint"
+}
+
+apply from: "$SUB_PROJECT"
+apply from: "$DOC_CREATION"
+apply plugin: 'kotlinx-atomicfu'
+
+dependencies {
+    compileOnly "org.jetbrains.kotlin:kotlin-stdlib-jdk7:$KOTLIN_VERSION"
+    testRuntimeOnly "org.junit.vintage:junit-vintage-engine:$JUNIT_VINTAGE_VERSION"
+    testCompileOnly "io.kotlintest:kotlintest-runner-junit5:$KOTLIN_TEST_VERSION", excludeArrow
+    testImplementation project(":arrow-core-test")
+}

arrow-libs/core/arrow-continuations/gradle.properties

@@ -0,0 +1,6 @@
+# Maven publishing configuration
+POM_NAME=Arrow Continuations
+POM_ARTIFACT_ID=arrow-continuations
+POM_PACKAGING=jar
+# Build configuration
+kapt.incremental.apt=false

arrow-libs/core/arrow-continuations/src/main/kotlin/arrow/continuations/generic/DelimContScope.kt

@@ -0,0 +1,121 @@
+package arrow.continuations.generic
+
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.loop
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
+import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn
+import kotlin.coroutines.resume
+import kotlin.coroutines.suspendCoroutine
+
+/**
+ * Implements delimited continuations with with no multi shot support (apart from shiftCPS which trivially supports it).
+ *
+ * For a version that simulates multishot (albeit with drawbacks) see [MultiShotDelimContScope].
+ * For a version that allows nesting [reset] and calling parent scopes inside inner scopes see [NestedDelimContScope].
+ *
+ * The basic concept here is appending callbacks and polling for a result.
+ * Every shift is evaluated until it either finishes (short-circuit) or suspends (called continuation). When it suspends its
+ *  continuation is appended to a list waiting to be invoked with the final result of the block.
+ * When running a function we jump back and forth between the main function and every function inside shift via their continuations.
+ */
+class DelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedScope<R> {
+
+  /**
+   * Variable used for polling the result after suspension happened.
+   */
+  private val resultVar = atomic<R?>(null)
+
+  /**
+   * Variable for the next shift block to (partially) run, if it is empty that usually means we are done.
+   */
+  private val nextShift = atomic<(suspend () -> R)?>(null)
+
+  /**
+   * "Callbacks"/partially evaluated shift blocks which now wait for the final result
+   */
+  // TODO This can be append only, but needs fast reversed access
+  private val shiftFnContinuations = mutableListOf<Continuation<R>>()
+
+  /**
+   * Small wrapper that handles invoking the correct continuations and appending continuations from shift blocks
+   */
+  data class SingleShotCont<A, R>(
+    private val continuation: Continuation<A>,
+    private val shiftFnContinuations: MutableList<Continuation<R>>
+  ) : DelimitedContinuation<A, R> {
+    override suspend fun invoke(a: A): R = suspendCoroutine { resumeShift ->
+      shiftFnContinuations.add(resumeShift)
+      continuation.resume(a)
+    }
+  }
+
+  /**
+   * Wrapper that handles invoking manually cps transformed continuations
+   */
+  data class CPSCont<A, R>(
+    private val runFunc: suspend DelimitedScope<R>.(A) -> R
+  ) : DelimitedContinuation<A, R> {
+    override suspend fun invoke(a: A): R = DelimContScope<R> { runFunc(a) }.invoke()
+  }
+
+  /**
+   * Captures the continuation and set [func] with the continuation to be executed next by the runloop.
+   */
+  override suspend fun <A> shift(func: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
+    suspendCoroutine { continueMain ->
+      val delCont = SingleShotCont(continueMain, shiftFnContinuations)
+      assert(nextShift.compareAndSet(null, suspend { this.func(delCont) }))
+    }
+
+  /**
+   * Same as [shift] except we never resume execution because we only continue in [c].
+   */
+  override suspend fun <A, B> shiftCPS(func: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing =
+    suspendCoroutine {
+      assert(nextShift.compareAndSet(null, suspend { func(CPSCont(c)) }))
+    }
+
+  /**
+   * Unsafe if [f] calls [shift] on this scope! Use [NestedDelimContScope] instead if this is a problem.
+   */
+  override suspend fun <A> reset(f: suspend DelimitedScope<A>.() -> A): A =
+    DelimContScope(f).invoke()
+
+  fun invoke(): R {
+    f.startCoroutineUninterceptedOrReturn(this, Continuation(EmptyCoroutineContext) { result ->
+      resultVar.value = result.getOrThrow()
+    }).let {
+      if (it == COROUTINE_SUSPENDED) {
+        // we have a call to shift so we must start execution the blocks there
+        resultVar.loop { mRes ->
+          if (mRes == null) {
+            val nextShiftFn = nextShift.getAndSet(null)
+              ?: throw IllegalStateException("No further work to do but also no result!")
+            nextShiftFn.startCoroutineUninterceptedOrReturn(Continuation(EmptyCoroutineContext) { result ->
+              resultVar.value = result.getOrThrow()
+            }).let {
+              // If we suspended here we can just continue to loop because we should now have a new function to run
+              // If we did not suspend we short-circuited and are thus done with looping
+              if (it != COROUTINE_SUSPENDED) resultVar.value = it as R
+            }
+            // Break out of the infinite loop if we have a result
+          } else return@let
+        }
+      }
+      // we can return directly if we never suspended/called shift
+      else return@invoke it as R
+    }
+    assert(resultVar.value != null)
+    // We need to finish the partially evaluated shift blocks by passing them our result.
+    // This will update the result via the continuations that now finish up
+    for (c in shiftFnContinuations.asReversed()) c.resume(resultVar.value!!)
+    // Return the final result
+    return resultVar.value!!
+  }
+
+  companion object {
+    fun <R> reset(f: suspend DelimitedScope<R>.() -> R): R = DelimContScope(f).invoke()
+  }
+}

arrow-libs/core/arrow-continuations/src/main/kotlin/arrow/continuations/generic/DelimitedCont.kt

@@ -0,0 +1,34 @@
+package arrow.continuations.generic
+
+/**
+ * Base interface for a continuation
+ */
+interface DelimitedContinuation<A, R> {
+  suspend operator fun invoke(a: A): R
+}
+
+/**
+ * Base interface for our scope.
+ */
+// TODO This should be @RestrictSuspension but that breaks because a superclass is not considered to be correct scope
+// @RestrictsSuspension
+interface DelimitedScope<R> {
+  /**
+   * Capture the continuation and pass it to [f].
+   */
+  suspend fun <A> shift(f: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A
+
+  /**
+   * Manually cps transformed shift. This can be used to gain multishot without hacks, but it's not the nicest for a few reasons:
+   * - It does not inherit the scope, this means it will be hard to effects offering non-det to offer the same scope again...
+   * - it is manually cps transformed which means every helper between this and invoking the continuation also needs to be transformed.
+   */
+  suspend fun <A, B> shiftCPS(f: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing
+
+  /**
+   * Nest another scope inside the current one.
+   *
+   * It is important to use this over creating an unrelated scope because
+   */
+  suspend fun <A> reset(f: suspend DelimitedScope<A>.() -> A): A
+}

arrow-libs/core/arrow-continuations/src/main/kotlin/arrow/continuations/generic/MultiShotDelimCont.kt

@@ -0,0 +1,128 @@
+package arrow.continuations.generic
+
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.loop
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
+import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn
+import kotlin.coroutines.resume
+import kotlin.coroutines.suspendCoroutine
+
+/**
+ * (Simulated) Multishot capable delimited control scope
+ *
+ * This has several drawbacks:
+ * - f will rerun completely on multishot and only the results of [shift] are cached so any sideeffects outside of
+ *   [shift] will rerun!
+ * - This accumulates all results of [shift] (every argument passed when invoking the continuation) so on long running computations
+ *   this may keep quite a bit of memory
+ * - If the pure part before a multishot is expensive the multishot itself will have to rerun that, which makes it somewhat slow
+ * - This is terribly hard to implement properly with nested scopes (which this one does not support)
+ *
+ * As per usual understanding of [DelimContScope] is required as I will only be commenting differences for now.
+ */
+open class MultiShotDelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedScope<R> {
+
+  private val resultVar = atomic<R?>(null)
+  private val nextShift = atomic<(suspend () -> R)?>(null)
+
+  // TODO This can be append only and needs fast reversed access
+  private val shiftFnContinuations = mutableListOf<Continuation<R>>()
+
+  /**
+   * Keep the arguments passed to [DelimitedContinuation.invoke] to be able to replay the scope if necessary
+   */
+  // TODO This can be append only and needs fast random access and slicing
+  internal open val stack = mutableListOf<Any?>()
+
+  /**
+   * Our continuation now includes the function [f] to rerun on multishot, the current live (single-shot) continuation,
+   *  the current stack and the offset from that stack when this is created which is used to know when to resume normal
+   *  execution again on a replay.
+   */
+  class MultiShotCont<A, R>(
+    liveContinuation: Continuation<A>,
+    private val f: suspend DelimitedScope<R>.() -> R,
+    private val stack: MutableList<Any?>,
+    private val shiftFnContinuations: MutableList<Continuation<R>>
+  ) : DelimitedContinuation<A, R> {
+    // To make sure the continuation is only invoked once we put it in a nullable atomic and only access it through getAndSet
+    private val liveContinuation = atomic<Continuation<A>?>(liveContinuation)
+    private val stackOffset = stack.size
+
+    override suspend fun invoke(a: A): R =
+      when (val cont = liveContinuation.getAndSet(null)) {
+        // On multishot we replay with a prefilled stack from start to the point at which this object was created
+        //  (when the shift block this runs in was first called)
+        null -> PrefilledDelimContScope((stack.subList(0, stackOffset).toList() + a).toMutableList(), f).invoke()
+        // on the first pass we operate like a normal delimited scope but we also save the argument to the stack before resuming
+        else -> suspendCoroutine { resumeShift ->
+          shiftFnContinuations.add(resumeShift)
+          stack.add(a)
+          cont.resume(a)
+        }
+      }
+  }
+
+  data class CPSCont<A, R>(
+    private val runFunc: suspend DelimitedScope<R>.(A) -> R
+  ) : DelimitedContinuation<A, R> {
+    override suspend fun invoke(a: A): R = DelimContScope<R> { runFunc(a) }.invoke()
+  }
+
+  override suspend fun <A> shift(func: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
+    suspendCoroutine { continueMain ->
+      val c = MultiShotCont(continueMain, f, stack, shiftFnContinuations)
+      assert(nextShift.compareAndSet(null, suspend { this.func(c) }))
+    }
+
+  override suspend fun <A, B> shiftCPS(func: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing =
+    suspendCoroutine {
+      assert(nextShift.compareAndSet(null, suspend { func(CPSCont(c)) }))
+    }
+
+  // This assumes RestrictSuspension or at least assumes the user to never reference the parent scope in f.
+  override suspend fun <A> reset(f: suspend DelimitedScope<A>.() -> A): A =
+    MultiShotDelimContScope(f).invoke()
+
+  fun invoke(): R {
+    f.startCoroutineUninterceptedOrReturn(this, Continuation(EmptyCoroutineContext) { result ->
+      resultVar.value = result.getOrThrow()
+    }).let {
+      if (it == COROUTINE_SUSPENDED) {
+        resultVar.loop { mRes ->
+          if (mRes == null) {
+            val nextShiftFn = nextShift.getAndSet(null)
+              ?: throw IllegalStateException("No further work to do but also no result!")
+            nextShiftFn.startCoroutineUninterceptedOrReturn(Continuation(EmptyCoroutineContext) { result ->
+              resultVar.value = result.getOrThrow()
+            }).let {
+              if (it != COROUTINE_SUSPENDED) resultVar.value = it as R
+            }
+          } else return@let
+        }
+      } else return@invoke it as R
+    }
+    assert(resultVar.value != null)
+    for (c in shiftFnContinuations.asReversed()) c.resume(resultVar.value!!)
+    return resultVar.value!!
+  }
+
+  companion object {
+    fun <R> reset(f: suspend DelimitedScope<R>.() -> R): R = MultiShotDelimContScope(f).invoke()
+  }
+}
+
+class PrefilledDelimContScope<R>(
+  override val stack: MutableList<Any?>,
+  f: suspend DelimitedScope<R>.() -> R
+) : MultiShotDelimContScope<R>(f) {
+  var depth = 0
+
+  // Here we first check if we still have values in our local stack and if so we use those first
+  //  if not we delegate to the normal delimited control implementation
+  override suspend fun <A> shift(func: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
+    if (stack.size > depth) stack[depth++] as A
+    else super.shift(func).also { depth++ }
+}