IO rewrite ContinueOn impl (#1443)

Author: nomisRev

modules/effects/arrow-effects-data/src/main/kotlin/arrow/effects/IO.kt

@@ -206,6 +206,10 @@ sealed class IO<out A> : IOOf<A> {
     override fun unsafeRunTimedTotal(limit: Duration): Option<A> = unsafeResync(this, limit)
   }
 
+  internal data class Effect<out A>(val ctx: CoroutineContext, val effect: suspend () -> A) : IO<A>() {
+    override fun unsafeRunTimedTotal(limit: Duration): Option<A> = unsafeResync(this, limit)
+  }
+
   internal data class Bind<E, out A>(val cont: IO<E>, val g: (E) -> IO<A>) : IO<A>() {
     override fun unsafeRunTimedTotal(limit: Duration): Option<A> = throw AssertionError("Unreachable")
   }

modules/effects/arrow-effects-data/src/main/kotlin/arrow/effects/IORunLoop.kt

@@ -1,12 +1,12 @@
 package arrow.effects
 
-import arrow.core.Continuation
 import arrow.core.Either
 import arrow.core.Left
 import arrow.core.NonFatal
 import arrow.core.Right
 import arrow.effects.internal.Platform.ArrayStack
 import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.startCoroutine
 
 private typealias Current = IOOf<Any?>
@@ -71,7 +71,10 @@ internal object IORunLoop {
         }
         is IO.Async -> {
           // Return case for Async operations
-          return suspendInAsync(currentIO as IO<A>, bFirst, bRest, currentIO.k)
+          return suspendInAsync(currentIO as IO<A>, bFirst, bRest)
+        }
+        is IO.Effect -> {
+          return suspendInAsync(currentIO as IO<A>, bFirst, bRest)
         }
         is IO.Bind<*, *> -> {
           if (bFirst != null) {
@@ -87,15 +90,13 @@ internal object IORunLoop {
             bRest.push(bFirst)
           }
           val localCurrent = currentIO
-
           val currentCC = localCurrent.cc
-
           val localCont = currentIO.cont
 
           bFirst = { c: Any? -> IO.just(c) }
 
-          currentIO = IO.async { conn, cc ->
-            loop(localCont, conn, cc.asyncCallback(currentCC), null, null, null)
+          currentIO = IO.Bind(localCont) { a ->
+            IO.Effect(currentCC) { a }
           }
         }
         is IO.Map<*, *> -> {
@@ -136,20 +137,15 @@ internal object IORunLoop {
   private fun <A> suspendInAsync(
     currentIO: IO<A>,
     bFirst: BindF?,
-    bRest: CallStack?,
-    register: IOProc<Any?>
+    bRest: CallStack?
   ): IO<A> =
-  // Hitting an async boundary means we have to stop, however
-  // if we had previous `flatMap` operations then we need to resume
-  // the loop with the collected stack
-    when {
-      bFirst != null || (bRest != null && bRest.isNotEmpty()) ->
-        IO.Async { conn, cb ->
-          val rcb = RestartCallback(conn, cb as Callback)
-          rcb.prepare(bFirst, bRest)
-          register(conn, rcb)
-        }
-      else -> currentIO
+    // Hitting an async boundary means we have to stop, however if we had previous `flatMap` operations then we need to resume the loop with the collected stack
+    if (bFirst != null || (bRest != null && bRest.isNotEmpty())) {
+      IO.Async { conn, cb ->
+        loop(currentIO, conn, cb as Callback, null, bFirst, bRest)
+      }
+    } else {
+      currentIO
     }
 
   private fun loop(
@@ -216,9 +212,18 @@ internal object IORunLoop {
           if (rcb == null) {
             rcb = RestartCallback(conn, cb)
           }
-          rcb.prepare(bFirst, bRest)
+
+          // Return case for Async operations
+          rcb.start(currentIO, bFirst, bRest)
+          return
+        }
+        is IO.Effect -> {
+          if (rcb == null) {
+            rcb = RestartCallback(conn, cb)
+          }
+
           // Return case for Async operations
-          currentIO.k(conn, rcb)
+          rcb.start(currentIO, bFirst, bRest)
           return
         }
         is IO.Bind<*, *> -> {
@@ -235,15 +240,13 @@ internal object IORunLoop {
             bRest.push(bFirst)
           }
           val localCurrent = currentIO
-
           val currentCC = localCurrent.cc
-
           val localCont = currentIO.cont
 
           bFirst = { c: Any? -> IO.just(c) }
 
-          currentIO = IO.async { _, callback ->
-            loop(localCont, conn, callback.asyncCallback(currentCC), null, null, null)
+          currentIO = IO.Bind(localCont) { a ->
+            IO.Effect(currentCC) { a }
           }
         }
         is IO.Map<*, *> -> {
@@ -353,7 +356,12 @@ internal object IORunLoop {
    * A `RestartCallback` gets created only once, per [startCancelable] (`unsafeRunAsync`) invocation, once an `Async`
    * state is hit, its job being to resume the loop after the boundary, but with the bind call-stack restored.
    */
-  private data class RestartCallback(val connInit: IOConnection, val cb: Callback) : Callback {
+  private data class RestartCallback(val connInit: IOConnection, val cb: Callback) : Callback, kotlin.coroutines.Continuation<Any?> {
+
+    // Nasty trick to re-use `Continuation` with different CC.
+    private var _context: CoroutineContext = EmptyCoroutineContext
+    override val context: CoroutineContext
+      get() = _context
 
     private var conn: IOConnection = connInit
     private var canCall = false
@@ -364,12 +372,23 @@ internal object IORunLoop {
       this.conn = conn
     }
 
-    fun prepare(bFirst: BindF?, bRest: CallStack?) {
+    private fun prepare(bFirst: BindF?, bRest: CallStack?) {
       canCall = true
       this.bFirst = bFirst
       this.bRest = bRest
     }
 
+    fun start(async: IO.Async<Any?>, bFirst: BindF?, bRest: CallStack?) {
+      prepare(bFirst, bRest)
+      async.k(conn, this)
+    }
+
+    fun start(effect: IO.Effect<Any?>, bFirst: BindF?, bRest: CallStack?) {
+      prepare(bFirst, bRest)
+      this._context = effect.ctx
+      effect.effect.startCoroutine(this)
+    }
+
     override operator fun invoke(either: Either<Throwable, Any?>) {
       if (canCall) {
         canCall = false
@@ -379,25 +398,17 @@ internal object IORunLoop {
         }
       }
     }
-  }
-
-  private fun <T> ((Either<Throwable, T>) -> Unit).asyncCallback(currentCC: CoroutineContext): (Either<Throwable, T>) -> Unit =
-    { result ->
-      val func: suspend () -> Unit = { this(result) }
-
-      val normalResume: Continuation<Unit> = object : Continuation<Unit> {
-        override val context: CoroutineContext = currentCC
 
-        override fun resume(value: Unit) {
-        }
-
-        override fun resumeWithException(exception: Throwable) {
-          this@asyncCallback(Either.left(exception))
-        }
+    override fun resumeWith(result: Result<Any?>) {
+      if (canCall) {
+        canCall = false
+        result.fold(
+          { a -> loop(IO.Pure(a), conn, cb, this, bFirst, bRest) },
+          { e -> loop(IO.RaiseError(e), conn, cb, this, bFirst, bRest) }
+        )
       }
-
-      func.startCoroutine(normalResume)
     }
+  }
 
   private class RestoreContext(
     val old: IOConnection,