Restricted & Suspended delimited scopes (#259)

* Makes DelimContScope#reset suspend

* Add suspending tests to TestSuite

* Move result to promise as well. Fixes suspension in shift before immediate return

* Update MultiShotCont with new impl

* Entry point api functions

* ktLintFormat

* Fix `suspended` in `FxLaws` doesn't actually suspend

* Re-add eager versions, fix `RestrictSuspension` in `nullable` and add suspend version

* Test that shows flow comprehension hanging

* Adds Stream test that also hangs

* Remove multi-shot test demos

* Effect interface and computation builder

* Rename Effect to computations

* Disable ktlint in effect file

* Correct suspension requires ControlThrowable

* Upgrade ktlint to allow fun interfaces

* Fix support for nested scopes, add tests for not leaking Arrow Fx & KotlinX Coroutines

* Removes NestedDelimContScope, makes Multishot private until we have a compiler plugin solution for ordered binds. Prepares DelimitedScope to be implemented by all and SuspendingComputation to replace DelimContScope if it's possible to implement its shift on it

* unfinished progress to share with Simon

* Add docs, revert some changes for RestrictSuspension implementation. Keep shiftCPS internally

* progress toward encoding effects with suspend an restricted delimited scopes

* tests passing

* remove runRestrictedSuspension

* Fix ktlintMainSourceSetCheck issues

* Fix ktlintTestSourceSetCheck issues

* Fix ktlint issues for arrow-core-data

* Adapted Fx laws, test failing with suspension

* ktlint

* fix for Fx laws

* Remove BindSyntax leftover for nullable

* Remove old destructuring syntax for bind from Either doc

* Raquel suggestion

Co-authored-by: Rachel M. Carmena <[email protected]>

* Deprecate `validated` block in favor of `either` block.

* Deprecate Validated fx

* Either.toValidated/Nel + fix to Validated docs

* ktlint fixes

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

Author: 3 people

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

@@ -12,4 +12,5 @@ dependencies {
     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")
+    testImplementation "io.arrow-kt:arrow-fx-coroutines:$VERSION_NAME"
 }

arrow-libs/core/arrow-continuations/src/main/kotlin/arrow/continuations/Effect.kt

@@ -0,0 +1,23 @@
+package arrow.continuations
+
+import arrow.continuations.generic.DelimitedScope
+
+fun interface Effect<F> {
+  fun control(): DelimitedScope<F>
+
+  companion object {
+    suspend inline fun <Eff : Effect<*>, F, A> suspended(
+      crossinline eff: (DelimitedScope<F>) -> Eff,
+      crossinline just: (A) -> F,
+      crossinline f: suspend Eff.() -> A,
+    ): F =
+      Reset.suspended { just(f(eff(this))) }
+
+    inline fun <Eff : Effect<*>, F, A> restricted(
+      crossinline eff: (DelimitedScope<F>) -> Eff,
+      crossinline just: (A) -> F,
+      crossinline f: suspend Eff.() -> A,
+    ): F =
+      Reset.restricted { just(f(eff(this))) }
+  }
+}

arrow-libs/core/arrow-continuations/src/main/kotlin/arrow/continuations/Reset.kt

@@ -0,0 +1,38 @@
+package arrow.continuations
+
+import arrow.continuations.generic.ShortCircuit
+import arrow.continuations.generic.ControlThrowable
+import arrow.continuations.generic.DelimContScope
+import arrow.continuations.generic.RestrictedScope
+import arrow.continuations.generic.SuspendMonadContinuation
+import arrow.continuations.generic.SuspendedScope
+import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
+
+@PublishedApi
+internal object Reset {
+  /**
+   * Allows for building suspending single-shot computation blocks.
+   * For short-circuiting, or shifting, a [ShortCircuit] [ControlThrowable] is used.
+   * This ensures that any concurrent nested scopes are correctly closed.
+   *
+   * The usage of `try { ... } catch(e: Throwable) { ... }` will catch the [ShortCircuit] error,
+   * and will lead to recover of short-circuiting.
+   * You should always prefer to catch the most specific exception class, or
+   * use `Either.catch`, `Validated.catch` etc or `e.nonFatalOrThrow()`
+   * to ensure you're not catching `ShortCircuit`.
+   */
+  suspend fun <A> suspended(block: suspend SuspendedScope<A>.() -> A): A =
+    suspendCoroutineUninterceptedOrReturn { cont ->
+      SuspendMonadContinuation(cont, block)
+        .startCoroutineUninterceptedOrReturn()
+    }
+
+  /**
+   * Allows for building eager single-shot computation blocks.
+   * For short-circuiting, or shifting, `@RestrictSuspension` state machine is used.
+   * This doesn't allow nesting of computation blocks, or foreign suspension.
+   */
+  // TODO This should be @RestrictSuspension but that breaks because a superclass is not considered to be correct scope
+  fun <A> restricted(block: suspend RestrictedScope<A>.() -> A): A =
+    DelimContScope(block).invoke()
+}

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

@@ -0,0 +1,11 @@
+package arrow.continuations.generic
+
+/**
+ * A [Throwable] class intended for control flow.
+ * Instance of [ControlThrowable] should **not** be caught,
+ * and `arrow.core.NonFatal` does not catch this [Throwable].
+ * Thus by extension `Either.catch` and `Validated.catch` also don't catch [ControlThrowable].
+ */
+open class ControlThrowable : Throwable() {
+  override fun fillInStackTrace(): Throwable = this
+}

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

@@ -1,13 +1,11 @@
 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.intrinsics.suspendCoroutineUninterceptedOrReturn
 import kotlin.coroutines.resume
-import kotlin.coroutines.suspendCoroutine
 
 /**
  * Implements delimited continuations with with no multi shot support (apart from shiftCPS which trivially supports it).
@@ -20,17 +18,17 @@ import kotlin.coroutines.suspendCoroutine
  *  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> {
+internal open class DelimContScope<R>(private val f: suspend RestrictedScope<R>.() -> R) : RestrictedScope<R> {
 
   /**
    * Variable used for polling the result after suspension happened.
    */
-  private val resultVar = atomic<Any?>(EMPTY_VALUE)
+  private var resultVar: Any? = EMPTY_VALUE
 
   /**
    * 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)
+  private var nextShift: (suspend () -> R)? = null
 
   /**
    * "Callbacks"/partially evaluated shift blocks which now wait for the final result
@@ -45,9 +43,10 @@ class DelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedSco
     private val continuation: Continuation<A>,
     private val shiftFnContinuations: MutableList<Continuation<R>>
   ) : DelimitedContinuation<A, R> {
-    override suspend fun invoke(a: A): R = suspendCoroutine { resumeShift ->
+    override suspend fun invoke(a: A): R = suspendCoroutineUninterceptedOrReturn { resumeShift ->
       shiftFnContinuations.add(resumeShift)
       continuation.resume(a)
+      COROUTINE_SUSPENDED
     }
   }
 
@@ -63,43 +62,47 @@ class DelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedSco
   /**
    * Captures the continuation and set [f] with the continuation to be executed next by the runloop.
    */
-  override suspend fun <A> shift(f: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
-    suspendCoroutine { continueMain ->
+  override suspend fun <A> shift(f: suspend RestrictedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
+    suspendCoroutineUninterceptedOrReturn { continueMain ->
       val delCont = SingleShotCont(continueMain, shiftFnContinuations)
-      assert(nextShift.compareAndSet(null, suspend { this.f(delCont) }))
+      assert(nextShift == null)
+      nextShift = suspend { this.f(delCont) }
+      COROUTINE_SUSPENDED
     }
 
   /**
    * Same as [shift] except we never resume execution because we only continue in [c].
    */
-  override suspend fun <A, B> shiftCPS(f: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing =
-    suspendCoroutine {
-      assert(nextShift.compareAndSet(null, suspend { f(CPSCont(c)) }))
+  suspend fun <A, B> shiftCPS(f: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing =
+    suspendCoroutineUninterceptedOrReturn {
+      assert(nextShift == null)
+      nextShift = suspend { f(CPSCont(c)) }
+      COROUTINE_SUSPENDED
     }
 
   /**
    * 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 =
+  fun <A> reset(f: suspend DelimitedScope<A>.() -> A): A =
     DelimContScope(f).invoke()
 
   @Suppress("UNCHECKED_CAST")
   fun invoke(): R {
     f.startCoroutineUninterceptedOrReturn(this, Continuation(EmptyCoroutineContext) { result ->
-      resultVar.value = result.getOrThrow()
+      resultVar = 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 === EMPTY_VALUE) {
-            val nextShiftFn = nextShift.getAndSet(null)
-              ?: throw IllegalStateException("No further work to do but also no result!")
+        while (true) {
+          if (resultVar === EMPTY_VALUE) {
+            val nextShiftFn = requireNotNull(nextShift) { "No further work to do but also no result!" }
+            nextShift = null
             nextShiftFn.startCoroutineUninterceptedOrReturn(Continuation(EmptyCoroutineContext) { result ->
-              resultVar.value = result.getOrThrow()
+              resultVar = result.getOrThrow()
             }).let { nextRes ->
               // 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 (nextRes != COROUTINE_SUSPENDED) resultVar.value = nextRes as R
+              if (nextRes != COROUTINE_SUSPENDED) resultVar = nextRes as R
             }
             // Break out of the infinite loop if we have a result
           } else return@let
@@ -108,17 +111,15 @@ class DelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedSco
       // we can return directly if we never suspended/called shift
       else return@invoke it as R
     }
-    assert(resultVar.value !== EMPTY_VALUE)
+    assert(resultVar !== EMPTY_VALUE)
     // 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 as R)
+    for (c in shiftFnContinuations.asReversed()) c.resume(resultVar as R)
     // Return the final result
-    return resultVar.value as R
+    return resultVar as R
   }
 
   companion object {
-    fun <R> reset(f: suspend DelimitedScope<R>.() -> R): R = DelimContScope(f).invoke()
-
     @Suppress("ClassName")
     private object EMPTY_VALUE
   }

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

@@ -10,25 +10,21 @@ interface DelimitedContinuation<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.
+   * Exit the [DelimitedScope] with [R]
    */
-  suspend fun <A, B> shiftCPS(f: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing
+  suspend fun <A> shift(r: R): A
+}
 
+interface RestrictedScope<R> : DelimitedScope<R> {
   /**
-   * Nest another scope inside the current one.
-   *
-   * It is important to use this over creating an unrelated scope because
+   * Capture the continuation and pass it to [f].
    */
-  suspend fun <A> reset(f: suspend DelimitedScope<A>.() -> A): A
+  suspend fun <A> shift(f: suspend RestrictedScope<R>.(DelimitedContinuation<A, R>) -> R): A
+
+  override suspend fun <A> shift(r: R): A = shift { r }
 }
+
+interface SuspendedScope<R> : DelimitedScope<R>

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

@@ -6,6 +6,7 @@ import kotlin.coroutines.Continuation
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
 import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn
+import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
 import kotlin.coroutines.resume
 import kotlin.coroutines.suspendCoroutine
 
@@ -22,10 +23,11 @@ import kotlin.coroutines.suspendCoroutine
  *
  * 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> {
+internal open class MultiShotDelimContScope<R>(val f: suspend RestrictedScope<R>.() -> R) : RestrictedScope<R> {
 
+  // TODO Since runs blocking these don't need to be atomic
   private val resultVar = atomic<R?>(null)
-  private val nextShift = atomic<(suspend () -> R)?>(null)
+  private val nextShift = atomic<(suspend RestrictedScope<R>.() -> R)?>(null)
 
   // TODO This can be append only and needs fast reversed access
   private val shiftFnContinuations = mutableListOf<Continuation<R>>()
@@ -43,7 +45,7 @@ open class MultiShotDelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R)
    */
   class MultiShotCont<A, R>(
     liveContinuation: Continuation<A>,
-    private val f: suspend DelimitedScope<R>.() -> R,
+    private val f: suspend RestrictedScope<R>.() -> R,
     private val stack: MutableList<Any?>,
     private val shiftFnContinuations: MutableList<Continuation<R>>
   ) : DelimitedContinuation<A, R> {
@@ -71,19 +73,22 @@ open class MultiShotDelimContScope<R>(val f: suspend DelimitedScope<R>.() -> 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 ->
+  override suspend fun <A> shift(func: suspend RestrictedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
+    suspendCoroutineUninterceptedOrReturn { continueMain ->
       val c = MultiShotCont(continueMain, f, stack, shiftFnContinuations)
-      assert(nextShift.compareAndSet(null, suspend { this.func(c) }))
+      val s: suspend RestrictedScope<R>.() -> R = { this.func(c) }
+      assert(nextShift.compareAndSet(null, s))
+      COROUTINE_SUSPENDED
     }
 
-  override suspend fun <A, B> shiftCPS(func: suspend (DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing =
+  suspend fun <A, B> shiftCPS(func: suspend DelimitedScope<R>.(DelimitedContinuation<A, B>) -> R, c: suspend DelimitedScope<B>.(A) -> B): Nothing =
     suspendCoroutine {
-      assert(nextShift.compareAndSet(null, suspend { func(CPSCont(c)) }))
+      val s: suspend DelimitedScope<R>.() -> R = { func(CPSCont(c)) }
+      assert(nextShift.compareAndSet(null, s))
     }
 
   // 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 =
+  suspend fun <A> reset(f: suspend DelimitedScope<A>.() -> A): A =
     MultiShotDelimContScope(f).invoke()
 
   fun invoke(): R {
@@ -95,7 +100,7 @@ open class MultiShotDelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R)
           if (mRes == null) {
             val nextShiftFn = nextShift.getAndSet(null)
               ?: throw IllegalStateException("No further work to do but also no result!")
-            nextShiftFn.startCoroutineUninterceptedOrReturn(Continuation(EmptyCoroutineContext) { result ->
+            nextShiftFn.startCoroutineUninterceptedOrReturn(this, Continuation(EmptyCoroutineContext) { result ->
               resultVar.value = result.getOrThrow()
             }).let {
               if (it != COROUTINE_SUSPENDED) resultVar.value = it as R
@@ -110,19 +115,21 @@ open class MultiShotDelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R)
   }
 
   companion object {
-    fun <R> reset(f: suspend DelimitedScope<R>.() -> R): R = MultiShotDelimContScope(f).invoke()
+    internal fun <R> reset(f: suspend RestrictedScope<R>.() -> R): R = MultiShotDelimContScope(f).invoke()
   }
 }
 
-class PrefilledDelimContScope<R>(
+private class PrefilledDelimContScope<R>(
   override val stack: MutableList<Any?>,
-  f: suspend DelimitedScope<R>.() -> R
+  f: suspend RestrictedScope<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 =
+  override suspend fun <A> shift(func: suspend RestrictedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
     if (stack.size > depth) stack[depth++] as A
-    else super.shift(func).also { depth++ }
+    else
+      @Suppress("ILLEGAL_RESTRICTED_SUSPENDING_FUNCTION_CALL")
+      super.shift(func).also { depth++ }
 }