Provide nullable continuation based on DelimitedScope

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

@@ -25,7 +25,7 @@ class DelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedSco
   /**
    * Variable used for polling the result after suspension happened.
    */
-  private val resultVar = atomic<R?>(null)
+  private val resultVar = atomic<Any?>(EMPTY_VALUE)
 
   /**
    * Variable for the next shift block to (partially) run, if it is empty that usually means we are done.
@@ -61,20 +61,20 @@ class DelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedSco
   }
 
   /**
-   * Captures the continuation and set [func] with the continuation to be executed next by the runloop.
+   * Captures the continuation and set [f] with the continuation to be executed next by the runloop.
    */
-  override suspend fun <A> shift(func: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
+  override suspend fun <A> shift(f: suspend DelimitedScope<R>.(DelimitedContinuation<A, R>) -> R): A =
     suspendCoroutine { continueMain ->
       val delCont = SingleShotCont(continueMain, shiftFnContinuations)
-      assert(nextShift.compareAndSet(null, suspend { this.func(delCont) }))
+      assert(nextShift.compareAndSet(null, suspend { this.f(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 =
+  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 { func(CPSCont(c)) }))
+      assert(nextShift.compareAndSet(null, suspend { f(CPSCont(c)) }))
     }
 
   /**
@@ -83,22 +83,23 @@ class DelimContScope<R>(val f: suspend DelimitedScope<R>.() -> R) : DelimitedSco
   override suspend 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()
     }).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) {
+          if (mRes === EMPTY_VALUE) {
             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 {
+            }).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 (it != COROUTINE_SUSPENDED) resultVar.value = it as R
+              if (nextRes != COROUTINE_SUSPENDED) resultVar.value = nextRes as R
             }
             // Break out of the infinite loop if we have a result
           } else return@let
@@ -107,15 +108,18 @@ 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 != null)
+    assert(resultVar.value !== 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!!)
+    for (c in shiftFnContinuations.asReversed()) c.resume(resultVar.value as R)
     // Return the final result
-    return resultVar.value!!
+    return resultVar.value as R
   }
 
   companion object {
     fun <R> reset(f: suspend DelimitedScope<R>.() -> R): R = DelimContScope(f).invoke()
+
+    @Suppress("ClassName")
+    private object EMPTY_VALUE
   }
 }

arrow-core-data/src/main/kotlin/arrow/core/computations/nullable.kt

@@ -0,0 +1,18 @@
+package arrow.core.computations
+
+import arrow.continuations.generic.DelimContScope
+import arrow.continuations.generic.DelimitedScope
+import arrow.core.computations.suspended.BindSyntax
+
+@Suppress("ClassName")
+object nullable {
+  operator fun <A> invoke(func: suspend BindSyntax.() -> A?): A? =
+    DelimContScope.reset { NullableBindSyntax(this).func() }
+
+  private class NullableBindSyntax<R>(
+    scope: DelimitedScope<R?>
+  ) : BindSyntax, DelimitedScope<R?> by scope {
+    override suspend fun <A> A?.invoke(): A =
+      this ?: shift { null }
+  }
+}

arrow-core-data/src/main/kotlin/arrow/core/computations/suspended/BindSyntax.kt

@@ -0,0 +1,17 @@
+package arrow.core.computations.suspended
+
+/**
+ * Running A? in the context of [nullable]
+ *
+ * ```
+ * nullable {
+ *   val one = 1.invoke() // using invoke
+ *   val bigger = (one.takeIf{ it > 1 }).invoke() // using invoke on expression
+ *   bigger
+ * }
+ * ```
+ */
+// TODO: this will become interface fun when they support suspend in the next release
+interface BindSyntax {
+  suspend operator fun <A> A?.invoke(): A
+}

arrow-core-data/src/test/kotlin/arrow/core/computations/NullableTest.kt

@@ -0,0 +1,61 @@
+package arrow.core.computations
+
+import arrow.core.test.UnitSpec
+import io.kotlintest.shouldBe
+
+class NullableTest : UnitSpec() {
+
+  init {
+    "simple case" {
+      nullable {
+        "s".length().invoke()
+      } shouldBe 1
+    }
+    "multiple types" {
+      nullable {
+        val number = "s".length()
+        val string = number.toString()()
+        string
+      } shouldBe "1"
+    }
+    "short circuit" {
+      nullable {
+        val number: Int = "s".length()
+        (number.takeIf { it > 1 }?.toString())()
+        throw IllegalStateException("This should not be executed")
+      } shouldBe null
+    }
+    "when expression" {
+      nullable {
+        val number = "s".length()
+        val string = when (number) {
+          1 -> number.toString()
+          else -> null
+        }.invoke()
+        string
+      } shouldBe "1"
+    }
+    "if expression" {
+      nullable {
+        val number = "s".length()
+        val string = if (number == 1) {
+          number.toString()
+        } else {
+          null
+        }.invoke()
+        string
+      } shouldBe "1"
+    }
+    "if expression short circuit" {
+      nullable {
+        val number = "s".length()
+        val string = if (number != 1) {
+          number.toString()
+        } else {
+          null
+        }()
+        string
+      } shouldBe null
+    }
+  }
+}