使用 Kotlin 协程实现 yield 函数

914 阅读1分钟
fun main() {
    val block: () -> Unit = YieldBlock {
        yield()         // #1
        println(111)    // #2 - co1 of #1
        println(222)    // #3 |
        yield()         // #4 |
        println(333)    // #5 |           -- co4 of #4
        println(444)    // #6 - co1 of #1 -- co4 of #4
    }

    block.invoke() // #1
    block.invoke() // #2#3#4
}

interface YieldScope {
    suspend fun yield()
}

class YieldBlock(
    block: suspend YieldScope.() -> Unit
) : () -> Unit, YieldScope, Continuation<Unit> {

    private var next: Continuation<Unit>? = block.createCoroutine(this, this)

    override fun invoke() {
        next?.resume(Unit)
    }

    override suspend fun yield() {
        suspendCoroutine<Unit> { co ->
            next = co // #7
        }
    }

    override fun resumeWith(result: Result<Unit>) { // #8
        next = null
        println("all end")
    }

    override val context: CoroutineContext
        get() = EmptyCoroutineContext
}

#2 - #6#1 的 continuation

#5 - #6#4 的 continuation

第一次执行 invoke() 时,运行至 #1 遇到 yield()yield() 的实现 #7 未触发后续 co1 的 resume() 方法,所以后续 co1 被挂起暂停不会继续执行,缓存 co1 等待 resume() 被调用恢复执行

第二次执行 invoke() 时,co1 被恢复执行,运行至 #4 遇到 yield() 被挂起暂停,并缓存后续 co4 等待恢复

#1 - #6 全部执行完毕,才会执行 #8

如果替换 #7 处为 co.resume(Unit) ,则仅需要 block.invoke() 一次,全部代码块都会被执行,而不会出现挂起暂停