Using Kotlin to implemant python style coroutine. 用 Kotlin 实现 Python 风格的协程 API

Generator.kt

import kotlin.coroutines.*

fun <T> generator(block: suspend GeneratorScope<T>.() -> Unit) = Generator(block)

class Generator<T>(private val block: suspend GeneratorScope<T>.() -> Unit) {
    operator fun iterator(): Iterator<T> = GeneratorIterator(block)
}

@RestrictsSuspension
interface GeneratorScope<T> {
    suspend fun yield(value: T)
}

sealed class State {
    class NotReady(val continuation: Continuation<Unit>) : State()
    class Ready<T>(val continuation: Continuation<Unit>, val v: T) : State()
    object Done : State()
}

class GeneratorIterator<T>(
    private val block: suspend GeneratorScope<T>.() -> Unit
) : Iterator<T>, GeneratorScope<T>, Continuation<Unit> {
    override val context: CoroutineContext
        get() = EmptyCoroutineContext

    private var state: State

    init {
        val coroutine = block.createCoroutine(this, this)
        state = State.NotReady(coroutine)
    }

    private fun resume() {
        when (val s = state) {
            is State.NotReady -> s.continuation.resume(Unit)
            else -> { /* do nothing */
            }
        }
    }

    override fun hasNext(): Boolean {
        resume()
        return state != State.Done
    }

    override fun next(): T = when (val s = state) {
        is State.NotReady -> {
            resume()
            next()
        }
        is State.Ready<*> -> {
            state = State.NotReady(s.continuation)
            (s as State.Ready<T>).v
        }
        State.Done -> throw IndexOutOfBoundsException("No value left.")
    }

    override suspend fun yield(value: T) {
        suspendCoroutine<Unit> { continuation ->
            state = when (state) {
                is State.NotReady -> State.Ready(continuation, value)
                is State.Ready<*> -> throw IllegalStateException("Cannot yield a value while ready.")
                is State.Done -> throw IllegalStateException("Cannot yield a value while done.")
            }
        }
    }

    /** Coroutine completion callback 协程完成回调 */
    override fun resumeWith(result: Result<Unit>) {
        state = State.Done
        result.getOrThrow()
    }
}

/** Usage 使用 */
fun main() {
    val numsGen = generator<Int> {
        yield(1)
        yield(2)
    }

    for (i in numsGen) {
        println("[$i]")
    }
}

See https://chenhe.me/post/dive-into-kotlin-coroutines-2