Kotlin first-class inline functions

decorators.kt

import kotlin.coroutines.Continuation
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.coroutines.RestrictsSuspension
import kotlin.coroutines.cancellation.CancellationException
import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn
import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn

private fun main() {
  val d: Decorator<Unit, Unit, Unit> = {
    test { iteration(it) }
  }
  d {
    return
  }
}

private inline fun test(block: (Unit) -> Unit) {
  repeat(3) {
    try {
      block(Unit)
    } finally {
      return@repeat
    }
  }
}

private inline fun calculate(transform: (Int) -> String): String =
  try {
    (1..5).map {
      var failed = true
      try {
        transform(it)
        failed = false
      } finally {
        if (failed) transform(it + 1)
      }
    }.joinToString()
  } finally {
    println("releasing resources")
  }

public typealias Decorator<T, R, S> = suspend DecoratorScope<T, R>.() -> S

public inline operator fun <T, R, S> Decorator<T, R, S>.invoke(block: (T) -> R): S =
  with(DecoratorCoroutine<T, R, S>()) {
    (beginDecorator(this@invoke) ?: provideIteration(block) {
      provideIteration(block) {
        provideIteration(block)
      }
    }).getOrThrow()
  }

@PublishedApi
internal inline fun <T, R, S> DecoratorCoroutine<T, R, S>.provideIteration(
  block: (T) -> R,
  handleFinally: () -> Result<S> = { TODO("Edgecase: calling `iteration` in `finally` that's caused by `iteration`") }
): Result<S> {
  var result: Result<S>? = null
  while (result == null) {
    var resultSet = false
    try {
      result = provideResult(Result.success(block(value)))
      resultSet = true
    } catch (e: Throwable) {
      result = provideResult(Result.failure(e))
      resultSet = true
    } finally {
      if (!resultSet) {
        val exc = CancellationException()
        result = provideResult(Result.failure(exc))
        if (result == null) result = handleFinally()
        if (result.exceptionOrNull() !== exc) continue
      }
    }
  }
  return result
}

/**
 * The [decorator] suspends for the iteration so that the one lambda can be run as
 * two separate parts, without needing to wrap an inlined block.
 */
@PublishedApi
internal class DecoratorCoroutine<T, R, S> : Continuation<S>, DecoratorScope<T, R> {
  private var continueAfterIteration: Continuation<R>? = null
  var value: T = null as T
    private set
  private var ultimateResult: Result<S>? = null

  override suspend fun iteration(param: T): R {
    check(continueAfterIteration == null) {
      "Iteration function somehow referenced itself!"
    }

    return suspendCoroutineUninterceptedOrReturn {
      continueAfterIteration = it
      value = param
      COROUTINE_SUSPENDED
    }
  }

  @PublishedApi
  internal fun beginDecorator(decorator: Decorator<T, R, S>): Result<S>? {
    check(ultimateResult == null) { "Decorator already completed" }

    return runCatching {
      val res = decorator.startCoroutineUninterceptedOrReturn(
        receiver = this,
        completion = this
      )
      if (res == COROUTINE_SUSPENDED) return null
      res as S
    }
  }

  @PublishedApi
  internal fun provideResult(value: Result<R>): Result<S>? {
    check(ultimateResult == null) { "Decorator already completed" }
    val cont = continueAfterIteration ?: error("Iteration not invoked")
    continueAfterIteration = null
    cont.resumeWith(value)
    return ultimateResult.also { ultimateResult = null }
  }

  override val context get() = EmptyCoroutineContext
  override fun resumeWith(result: Result<S>) {
    ultimateResult = result
  }
}

@RestrictsSuspension
public sealed interface DecoratorScope<T, R> {
  public suspend fun iteration(param: T): R
}