A little trampolining implementation I wrote to pr...

# coroutines

[Youssef Shoaib [MOD]]

A little trampolining implementation I wrote to provide functionality similar to DeepRecursiveFunction while allowing foreign suspension. Any critiques?

internal fun <T> Continuation<T>.resumeWithIntercepted(result: Result<T>) {
  context.trampoline.next { resumeWith(result) }
}

@OptIn(InternalCoroutinesApi::class)
internal fun CoroutineContext.withTrampoline(): CoroutineContext {
  val interceptor = this[ContinuationInterceptor].let {
    if (it is Trampoline) it.interceptor else it
  }
  return this + if (interceptor is Delay) TrampolineWithDelay(interceptor, interceptor) else Trampoline(interceptor)
}

@InternalCoroutinesApi
private class TrampolineWithDelay(interceptor: ContinuationInterceptor?, delay: Delay) :
  Trampoline(interceptor), Delay by delay

private open class Trampoline(val interceptor: ContinuationInterceptor?) :
  AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor {
  private var toRun: (() -> Unit)? = null
  fun next(block: () -> Unit) {
    check(toRun == null) { "Already running a block: $toRun" }
    toRun = block
  }

  override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> =
    TrampolineContinuation(continuation).let {
      interceptor?.interceptContinuation(it) ?: it
    }

  override fun releaseInterceptedContinuation(continuation: Continuation<*>) {
    interceptor?.releaseInterceptedContinuation(continuation)
  }

  private inner class TrampolineContinuation<T>(val cont: Continuation<T>) : Continuation<T> {
    override val context: CoroutineContext = cont.context

    override fun resumeWith(result: Result<T>) {
      cont.resumeWith(result)
      while (true) {
        (toRun ?: return).also { toRun = null }.invoke()
      }
    }
  }
}

private val CoroutineContext.trampoline: Trampoline
  get() =
    this[ContinuationInterceptor] as? Trampoline ?: error("No trampoline in context: $this")

Could likely be improved by defunctionalizing toRun lambdas if they're known to only need e.g. 2 fields.