使用Kotlin协程实现python的generator功能

一、要实现python的generator功能，首先要看下该generator是怎样的调用，如下的调用示例：

fun testGenerator() {
    val generator = generator<Int> {
        println("generator $it")
        yield(10)
        yield(22)
    }

    for (element in generator(10)) {
        println(element)
    }
}

 

从上面可以看出：

1、generator接收的是一个可挂起的函数：该函数是一个伴生函数，对应的类包含挂起函数yield

2、generator的返回值是一个函数：该函数有一个参数，返回值是支持迭代器的对象

3、要在该对象进行挂起和唤醒，则该对象支持协程

所以总结：所以返回的对象支持：迭代器、支持挂起函数yield，是一个协程

 

通过上面的总结，我们可以先大致构造出对应的generator如下所示

interface Generator<T> {
    operator fun iterator(): Iterator<T>
}

interface GeneratorScope<T> {
    suspend fun yield(parameter: T)
}

fun <T> generator(block: suspend GeneratorScope<T>.(T) -> Unit): (T) -> Generator<T> {
    return {
　　　　　// 返回的是一个接受类型为T的参数的函数
        parameter ->
        GeneratorImpl(block, parameter)
    }
}

class GeneratorImpl<T>(val block: suspend GeneratorScope<T>.(T) -> Unit, val parameter: T): Generator<T> {
    override fun iterator(): Iterator<T> {
        // 最终实际是调用该类，那么该类需要支持之前所有的特性：迭代、yield、协程
        return GeneratorIterator(block, parameter)
    }
}

 

二、通过一中的GeneratorIterator的备注，咱们可以构造该类如下

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

class GeneratorIterator<T>(block: suspend GeneratorScope<T>.(T) -> Unit, parameter: T): Iterator<T>, GeneratorScope<T>, Continuation<Any?> {
    override val context = EmptyCoroutineContext
    private var mState: State

    init {
        // 由于要协程开始调用block，并且接收的跟block是同一个类
        // 所以这里使用GeneratorScope的挂起伴生对象
        // 由于startBlock只是要启动协程，不要参数所以这里的伴生对象函数参数为空
        val startBlock: suspend GeneratorScope<T>.() -> Unit = {block(parameter)}
        // 正常使用继承block对应类的方式，所以这里传入的接收者是this
        // 因为该类使用协程，所以completion也是this
        val startCoroutine = startBlock.createCoroutine(this, this)
        // 所有的状态变化都是放在mState
        mState = State.NotReady(startCoroutine)
    }

。。。。。。。

 

如上的GeneratorIterator接收传入的要执行的block和初始要执行的参数，并且继承了Iterator, GeneratorScope, Continuation

 

三、完整的代码如下所示：

interface Generator<T> {
    operator fun iterator(): Iterator<T>
}

interface GeneratorScope<T> {
    suspend fun yield(parameter: T)
}

fun <T> generator(block: suspend GeneratorScope<T>.(T) -> Unit): (T) -> Generator<T> {
    return {
        parameter ->
        GeneratorImpl(block, parameter)
    }
}

class GeneratorImpl<T>(val block: suspend GeneratorScope<T>.(T) -> Unit, 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 parameter: T): State()
    object Done: State()
}

class GeneratorIterator<T>(block: suspend GeneratorScope<T>.(T) -> Unit, parameter: T): Iterator<T>, GeneratorScope<T>, Continuation<Any?> {
    override val context = EmptyCoroutineContext
    private var mState: State

    init {
        val startBlock: suspend GeneratorScope<T>.() -> Unit = {block(parameter)}
        val startCoroutine = startBlock.createCoroutine(this, this)
        mState = State.NotReady(startCoroutine)
    }

    override fun hasNext(): Boolean {
        resume()

        return mState != State.Done
    }

    override fun next(): T {
        return when(val currentState = mState) {
            is State.NotReady -> {
                resume()
                return next()
            }
            is State.Ready<*> -> {
                // 已经是ready之后，要将state转为notready，然后返回ready里面的value
                mState = State.NotReady(currentState.continuation)
                return (currentState as State.Ready<T>).parameter
            }

            is State.Done -> throw IllegalStateException("next have done error")
        }
    }

    // 设置了本身的complete之后，在suspend的block完成之后会调用该resumeWith，则需要将state设置为done
    override fun resumeWith(result: Result<Any?>) {
        mState = State.Done

        result.getOrThrow()
    }

    // yield的状态只能是从NotReady到Ready
    // suspendCoroutine挂起当前的函数
    // 会调用resume执行该block
    // 执行完block之后会挂起当前的协程
    override suspend fun yield(parameter: T) = suspendCoroutine {
        continuation ->
        mState = when(mState) {
            is State.NotReady -> State.Ready(continuation, parameter)
            is State.Ready<*> -> throw IllegalStateException("yield ready error")
            is State.Done -> throw IllegalStateException("yield done error")
        }
    }

    private fun resume() {
        when(val currentState = mState) {
            // 调用resume之后则会继续执行传入generator的block函数
            is State.NotReady -> currentState.continuation.resume(Unit)
            else -> Unit
        }
    }
}