diff --git a/libraries/stdlib/src/kotlin/coroutines/experimental/SequenceBuilder.kt b/libraries/stdlib/src/kotlin/coroutines/experimental/SequenceBuilder.kt index 3f1b16c9025..e8251d24fd7 100644 --- a/libraries/stdlib/src/kotlin/coroutines/experimental/SequenceBuilder.kt +++ b/libraries/stdlib/src/kotlin/coroutines/experimental/SequenceBuilder.kt @@ -94,10 +94,11 @@ public abstract class SequenceBuilder internal constructor() { private typealias State = Int private const val State_NotReady: State = 0 -private const val State_ManyReady: State = 1 -private const val State_Ready: State = 2 -private const val State_Done: State = 3 -private const val State_Failed: State = 4 +private const val State_ManyNotReady: State = 1 +private const val State_ManyReady: State = 2 +private const val State_Ready: State = 3 +private const val State_Done: State = 4 +private const val State_Failed: State = 5 private class SequenceBuilderIterator : SequenceBuilder(), Iterator, Continuation { private var state = State_NotReady @@ -109,10 +110,15 @@ private class SequenceBuilderIterator : SequenceBuilder(), Iterator, Co while (true) { when (state) { State_NotReady -> {} - State_ManyReady -> - if (nextIterator!!.hasNext()) return true else nextIterator = null + State_ManyNotReady -> + if (nextIterator!!.hasNext()) { + state = State_ManyReady + return true + } else { + nextIterator = null + } State_Done -> return false - State_Ready -> return true + State_Ready, State_ManyReady -> return true else -> throw exceptionalState() } @@ -125,8 +131,11 @@ private class SequenceBuilderIterator : SequenceBuilder(), Iterator, Co override fun next(): T { when (state) { - State_NotReady -> return nextNotReady() - State_ManyReady -> return nextIterator!!.next() + State_NotReady, State_ManyNotReady -> return nextNotReady() + State_ManyReady -> { + state = State_ManyNotReady + return nextIterator!!.next() + } State_Ready -> { state = State_NotReady val result = nextValue as T diff --git a/libraries/stdlib/test/coroutines/SequenceBuilderTest.kt b/libraries/stdlib/test/coroutines/SequenceBuilderTest.kt index 81f7f320fd7..3b43257cf5e 100644 --- a/libraries/stdlib/test/coroutines/SequenceBuilderTest.kt +++ b/libraries/stdlib/test/coroutines/SequenceBuilderTest.kt @@ -287,4 +287,19 @@ class SequenceBuilderTest { effects.toList() ) } + + @Test + fun testInfiniteYieldAll() { + val values = buildIterator { + while (true) { + yieldAll((1..5).map { it }) + } + } + + var sum = 0 + repeat(10) { + sum += values.next() //.also(::println) + } + assertEquals(30, sum) + } }