Kotlin协程实现懒序列

懒序列是啥？

不使用就不初始化，调用时才加载的序列就是懒序列。kotlin中，可以利用协程的挂起和恢复特性，实现懒序列。

期望达到的效果

fun main() {
    val arr = generator<Int> {
        for (i in 0..5) {
            println("$i 准备好啦。")
            yield(i)
        }
    }

    for (i in arr) {
        println("输出 $i")
    }
}

// 控制台输出
0 准备好啦。
输出 0
1 准备好啦。
输出 1
2 准备好啦。
输出 2
3 准备好啦。
输出 3
4 准备好啦。
输出 4
5 准备好啦。
输出 5

Process finished with exit code 0

难点分析

• 实现一个generator方法，接受一个函数，并返回一个序列
• 限定yield方法只能在generator方法的lambda表达式的作用域中使用
• generator方法接受的函数必须是挂起函数，否则无法达到懒加载的目的
• 为了支持 foreach 表达式，需要让generator方法返回一个Iterator对象

具体实现

import kotlin.coroutines.*

fun main() {
    val arr = generator<Int> {
        for (i in 0..5) {
            println("$i 准备好啦。")
            yield(i)
        }
    }

    for (i in arr) {
        println("输出 $i")
    }
}

fun <T> generator(block: suspend GeneratorScope<T>.() -> Unit): Iterator<T> {
    return GeneratorIterator(block)
}

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>(block: suspend GeneratorScope<T>.() -> Unit) :
    Iterator<T>, GeneratorScope<T>, Continuation<Unit> {
    override val context: CoroutineContext = EmptyCoroutineContext

    private var state: State

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

    /**
     * yield方法干了啥
     * 1.挂起当前方法，
     * 2.恢复遍历处的方法，并将值通过状态传给遍历处
     * 3.流转状态，将状态从 NotReady 改成 Ready
     */
    override suspend fun yield(value: T) = suspendCoroutine<Unit> {
        state = State.Ready(it, value)
    }

    /**
     * hasNext方法干了啥
     * 1.如果当前状态为 NotReady, 则恢复NotReady协程让状态进行到 Ready 或者 Done
     * 2.然后判断当前state是否是Done，来确定是否还有下一个元素
     */
    override fun hasNext(): Boolean {
        if (state is State.NotReady) {
            (state as State.NotReady).continuation.resume(Unit)
        }
        return state != State.Done
    }

    /**
     * next方法干了啥
     * 1.判断是否有下一个，如果有，则取出来返回，同时流转状态，Ready -> NotReady
     * 2.没有就报数组越界
     */
    override fun next(): T {
        if (hasNext()) {
            val currentState = state as State.Ready<T>
            state = State.NotReady(currentState.continuation)
            return currentState.nextValue
        }
        throw IndexOutOfBoundsException("No value left.")
    }

    override fun resumeWith(result: Result<Unit>) {
        state = State.Done
    }
}

/**
 * 限制yield方法只能在特定的lambda表达式中使用
 */
interface GeneratorScope<T> {
    suspend fun yield(value: T)
}