Created
July 20, 2022 08:35
-
-
Save ichenhe/64250d4b3f514acb0ea7f273759bf212 to your computer and use it in GitHub Desktop.
Using Kotlin to implemant python style coroutine. 用 Kotlin 实现 Python 风格的协程 API
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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]") | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
See https://chenhe.me/post/dive-into-kotlin-coroutines-2