kotlin实现Python Generator

使用kotlin协程实现Generator并支持传入初始参数

package leo.demo

import kotlin.coroutines.*

/**
 * 序列接口
 */
interface Generator<T> {
    operator fun iterator(): Iterator<T>
}

/**
 * 序列实现类
 */
class GeneratorImpl<T>
    (private val block: suspend GeneratorScope<T>.(T) -> Unit, private val parameter: T) :
    Generator<T> {
    override fun iterator(): Iterator<T> {
        return GeneratorIterator(block, parameter)
    }
}

/**
 * 序列状态机，表示是否准备好下一个值
 */
sealed class State {
    class NotReady(val continuation: Continuation<Unit>) : State()
    class Ready<T>(val continuation: Continuation<Unit>, val nextValue: T) : State()
    object Done : State()
}

/**
 * 序列迭代迭代器
 */
class GeneratorIterator<T>(private val block: suspend GeneratorScope<T>.(T) -> Unit, override val parameter: T) :
    GeneratorScope<T>(), Iterator<T>, Continuation<Any?> {
    override val context: CoroutineContext = EmptyCoroutineContext
    private var state: State

    init {
        //创建协程
        val coroutineBlock: suspend GeneratorScope<T>.() -> Unit = { block(parameter) }
        val star = coroutineBlock.createCoroutine(this, this)
        state = State.NotReady(star)
    }

    //挂起函数
    override suspend fun yield(value: T) = suspendCoroutine<Unit> { continuation: Continuation<Unit> ->
        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 java.lang.IllegalStateException("Cannot yield a value while done")
        }
    }

    //协程恢复
    private fun resume() {
        when (val currentState = state) {
            is State.NotReady -> currentState.continuation.resume(Unit)
        }
    }

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

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


    /**
     * 协程运行完毕后执行
     */
    override fun resumeWith(result: Result<Any?>) {
        state = State.Done
        result.getOrThrow();
    }
}

abstract class GeneratorScope<T> internal constructor() {
    protected abstract val parameter: T
    abstract suspend fun yield(value: T)
}

/**
 * 生成一个个序列生成器函数
 */
fun <T> generator(block: suspend GeneratorScope<T>.(T) -> Unit): (T) -> Generator<T> {
    return { parameter ->
        GeneratorImpl(block, parameter)
    }
}

fun main() {
    //生成一个个序列生成器函数
    val nums = generator<Int> { start ->
        for (i in 0..5) {
            yield(start + i)
        }
    }
    //幷成序列seq
    val seq = nums(10);

    //从序列获取值
    for (j in seq) {
        println("从seq获取值: $j")
    }
    // 官方提供了sequence
    val sequence=sequence{
        yield(1)
        yield(2)
        yieldAll(listOf(1,2,3,4))
    }
    for (x in sequence){
        println("x:$x")
    }
}