IO trampoline async (arrow-kt#1448)

Author: nomisRev

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

@@ -202,7 +202,7 @@ sealed class IO<out A> : IOOf<A> {
     override fun unsafeRunTimedTotal(limit: Duration): Option<A> = throw AssertionError("Unreachable")
   }
 
-  internal data class Async<out A>(val k: IOProc<A>) : IO<A>() {
+  internal data class Async<out A>(val shouldTrampoline: Boolean = false, val k: IOProc<A>) : IO<A>() {
     override fun unsafeRunTimedTotal(limit: Duration): Option<A> = unsafeResync(this, limit)
   }
 

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

@@ -6,6 +6,7 @@ import arrow.core.NonFatal
 import arrow.core.Right
 import arrow.core.nonFatalOrThrow
 import arrow.effects.internal.ArrowInternalException
+import arrow.effects.internal.Platform
 import arrow.effects.internal.Platform.ArrayStack
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
@@ -64,7 +65,7 @@ internal object IORunLoop {
             hasResult = true
             currentIO = null
           } catch (t: Throwable) {
-              currentIO = IO.RaiseError(t.nonFatalOrThrow())
+            currentIO = IO.RaiseError(t.nonFatalOrThrow())
           }
         }
         is IO.Async -> {
@@ -371,6 +372,12 @@ internal object IORunLoop {
     private var bFirst: BindF? = null
     private var bRest: CallStack? = null
 
+    private var contIndex: Int = 0
+    private var trampolineAfter: Boolean = false
+    private inline val shouldTrampoline inline get() = trampolineAfter || contIndex == Platform.maxStackDepthSize
+
+    private var value: IO<Any?>? = null
+
     fun contextSwitch(conn: IOConnection) {
       this.conn = conn
     }
@@ -379,10 +386,12 @@ internal object IORunLoop {
       canCall = true
       this.bFirst = bFirst
       this.bRest = bRest
+      contIndex++
     }
 
     fun start(async: IO.Async<Any?>, bFirst: BindF?, bRest: CallStack?) {
       prepare(bFirst, bRest)
+      trampolineAfter = async.shouldTrampoline
       async.k(conn, this)
     }
 
@@ -392,12 +401,30 @@ internal object IORunLoop {
       effect.effect.startCoroutine(this)
     }
 
+    private fun signal(result: IO<Any?>) {
+      // Allow GC to collect
+      val bFirst = this.bFirst
+      val bRest = this.bRest
+      this.bFirst = null
+      this.bRest = null
+      this._context = EmptyCoroutineContext
+
+      loop(result, conn, cb, this, bFirst, bRest)
+    }
+
     override operator fun invoke(either: Either<Throwable, Any?>) {
       if (canCall) {
         canCall = false
         when (either) {
-          is Either.Left -> loop(IO.RaiseError(either.a), conn, cb, this, bFirst, bRest)
-          is Either.Right -> loop(IO.Pure(either.b), conn, cb, this, bFirst, bRest)
+          is Either.Left -> IO.RaiseError(either.a)
+          is Either.Right -> IO.Pure(either.b)
+        }.let { r ->
+          if (shouldTrampoline) {
+            this.value = r
+            Platform.trampoline { trampoline() }
+          } else {
+            signal(r)
+          }
         }
       }
     }
@@ -406,11 +433,25 @@ internal object IORunLoop {
       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) }
-        )
+          { a -> IO.Pure(a) },
+          { e -> IO.RaiseError(e) }
+        ).let { r ->
+          if (shouldTrampoline) {
+            this.value = r
+            Platform.trampoline { trampoline() }
+          } else {
+            signal(r)
+          }
+        }
       }
     }
+
+    fun trampoline() {
+      val v = value
+      value = null
+      contIndex = 0
+      signal(v!!)
+    }
   }
 
   private class RestoreContext(

modules/effects/arrow-effects-data/src/main/kotlin/arrow/effects/internal/Utils.kt

@@ -162,9 +162,7 @@ object Platform {
       latch.tryAcquireSharedNanos(1, limit.nanoseconds)
     }
 
-    val eitherRef = ref
-
-    return when (eitherRef) {
+    return when (val eitherRef = ref) {
       null -> None
       is Either.Left -> throw eitherRef.a
       is Either.Right -> Some(eitherRef.b)
@@ -207,7 +205,7 @@ object Platform {
 
   @PublishedApi
   internal class TrampolineExecutor(val underlying: Executor) {
-    private var immediateQueue = Platform.ArrayStack<Runnable>()
+    private var immediateQueue = ArrayStack<Runnable>()
     @Volatile
     private var withinLoop = false
 
@@ -225,7 +223,7 @@ object Platform {
       else immediateQueue.push(runnable)
 
     private fun forkTheRest() {
-      class ResumeRun(val head: Runnable, val rest: Platform.ArrayStack<Runnable>) : Runnable {
+      class ResumeRun(val head: Runnable, val rest: ArrayStack<Runnable>) : Runnable {
         override fun run() {
           immediateQueue.pushAll(rest)
           immediateLoop(head)
@@ -235,7 +233,7 @@ object Platform {
       val head = immediateQueue.pop()
       if (head != null) {
         val rest = immediateQueue
-        immediateQueue = Platform.ArrayStack()
+        immediateQueue = ArrayStack()
         underlying.execute(ResumeRun(head, rest))
       }
     }

modules/effects/arrow-effects-data/src/test/kotlin/arrow/effects/IOTest.kt

@@ -6,8 +6,6 @@ import arrow.core.None
 import arrow.core.Right
 import arrow.core.Some
 import arrow.core.Tuple3
-import arrow.core.fix
-import arrow.core.flatMap
 import arrow.core.right
 import arrow.effects.IO.Companion.just
 import arrow.effects.extensions.io.async.async
@@ -284,7 +282,7 @@ class IOTest : UnitSpec() {
 
       val result =
         newSingleThreadContext("all").parMapN(
-          makePar(6), IO.just(1L).order(), makePar(4), IO.defer { IO.just(2L) }.order(), makePar(5), IO { 3L }.order()) { six, one, four, two, five, three -> listOf(six, one, four, two, five, three) }
+          makePar(6), just(1L).order(), makePar(4), IO.defer { just(2L) }.order(), makePar(5), IO { 3L }.order()) { six, one, four, two, five, three -> listOf(six, one, four, two, five, three) }
           .unsafeRunSync()
 
       result shouldBe listOf(6L, 1, 4, 2, 5, 3)
@@ -301,7 +299,7 @@ class IOTest : UnitSpec() {
 
       val result =
         newSingleThreadContext("all").parMapN(
-          makePar(6), IO.just(1L), makePar(4), IO.defer { IO.just(2L) }, makePar(5), IO { 3L }) { _, _, _, _, _, _ ->
+          makePar(6), just(1L), makePar(4), IO.defer { just(2L) }, makePar(5), IO { 3L }) { _, _, _, _, _, _ ->
           Thread.currentThread().name
         }.unsafeRunSync()
 
@@ -466,7 +464,20 @@ class IOTest : UnitSpec() {
           else just(ii)
         }
 
-      IO.just(1).flatMap { ioGuaranteeCase(0) }.unsafeRunSync() shouldBe size
+      just(1).flatMap { ioGuaranteeCase(0) }.unsafeRunSync() shouldBe size
+    }
+
+    "Async should be stack safe" {
+      val size = 5000
+
+      fun ioAsync(i: Int): IO<Int> = IO.async<Int> { _, cb ->
+        cb(Right(i))
+      }.flatMap { ii ->
+        if (ii < size) ioAsync(ii + 1)
+        else just(ii)
+      }
+
+      IO.just(1).flatMap(::ioAsync).unsafeRunSync() shouldBe size
     }
   }
 }