实现生成sequence序列的方法

kotlin

这段时间在学习kotlin的协程相关的内容,其中看到一个函数sequence(@BuilderInference block: suspend SequenceScope<T>.() -> Unit)，它可以一个接一个地或按任意大小的块生成序列(Sequences) 。这个函数
接受一个包含调用 yield() 和 yieldAll() 函数的 lambda 表达式。它们将一个元素返回给序列消费者并暂停
sequence()的执行，直到消费者请求下一个元素。 yield() 将单个元素作为参数；yieldAll()可以 接受一个 Iterable对象、一个Iterator或另一个Sequence。

sequence使用

官方sequence函数具体使用方法如下：

1
2
3
4
5
6
7
fun main() {
    val oddNumbers = sequence {
➡️       yield(1)
➡️       yieldAll(listOf(3, 5))
    }
    println(oddNumbers.toList())
}

注意到yield和yieldAll函数前面我用了➡️图标，这2个方法与普通的方法不太一样，是挂起方法，代码在编译器里会有 挂起标记。我这里用➡️图标替换。

运行上面这段程序，得到的结果是[1, 3, 5]。起初看到这段代码的时候觉得很神奇，协程还能这么用。我们
知道sequence是惰性序列，sequence生成的时候数据并没有产生，只有当sequence进行iterator操作的时候才会产生数据。

我们可以试着实现自己的sequence方法

完整实例代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
interface Generator<T> {
    operator fun iterator(): Iterator<T>
}

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

class GeneratorImpl<T>(private val block: suspend GeneratorScope<T>.(T) -> Unit, private val parameter: T) :
    Generator<T> {
    override fun iterator(): Iterator<T> = 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?> {
    var state: State

    init {
        val coroutineBlock: suspend GeneratorScope<T>.() -> Unit = {
            block(parameter)
        }
        val start = coroutineBlock.createCoroutine(this, this)
        state = State.NotReady(start)
    }

    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.")
            State.Done -> throw 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
            }
            State.Done -> throw IndexOutOfBoundsException("No value left")
        }
    }

    override val context: CoroutineContext
        get() = EmptyCoroutineContext


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

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

public fun <T> Generator<T>.toList(): List<T> {
    val destination = ArrayList<T>()
    for (item in this) {
        destination.add(item)
    }
    return when (destination.size) {
        0 -> emptyList<T>()
        1 -> listOf(destination[0])
        else -> destination
    }
}


public fun <T> Generator<T>.toList(): List<T> {
    val destination = ArrayList<T>()
    for (item in this) {
        destination.add(item)
    }
    return when (destination.size) {
        0 -> emptyList<T>()
        1 -> listOf(destination[0])
        else -> destination
    }
}

fun main() {
    val sequence = generator<Int> { p ->
        yield(p + 1)️
        yield(p + 2)
        yield(p + 3)
    }
    println(sequence(2).toList())
}