Replace generate implementaion

The new one is more efficient and contains `yieldAll`

Author: dzharkov

kotlinx-coroutines-generate/src/main/kotlin/generate.kt

@@ -1,42 +1,107 @@
 package kotlinx.coroutines
 
-import kotlin.coroutines.*
-import kotlin.coroutines.CoroutineIntrinsics.SUSPENDED
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.CoroutineIntrinsics
+import kotlin.coroutines.RestrictsSuspension
+import kotlin.coroutines.createCoroutine
 
 /**
- * Creates a Sequence object based on received coroutine [c].
- *
- * Each call of 'yield' suspend function within the coroutine lambda generates
- * next element of resulting sequence.
+ * Scope of [generate] block.
  */
-interface Generator<in T> {
-    suspend fun yield(value: T)
-}
+@RestrictsSuspension
+public abstract class Generator<in T> internal constructor() {
+    /**
+     * Yields a value in [generate] block.
+     */
+    public abstract suspend fun yield(value: T)
+
+    /**
+     * Yields potentially infinite sequence of iterator values in [generate] block.
+     */
+    public abstract suspend fun yieldAll(iterator: Iterator<T>)
+
+    /**
+     * Yields a collections of values in [generate] block.
+     */
+    public suspend fun yieldAll(elements: Iterable<T>) = yieldAll(elements.iterator())
 
-fun <T> generate(block: suspend Generator<T>.() -> Unit): Sequence<T> = GeneratedSequence(block)
+    /**
+     * Yields potentially infinite sequence of values in [generate] block.
+     */
+    public suspend fun yieldAll(sequence: Sequence<T>) = yieldAll(sequence.iterator())
+}
 
-private class GeneratedSequence<out T>(private val block: suspend Generator<T>.() -> Unit) : Sequence<T> {
-    override fun iterator(): Iterator<T> = GeneratedIterator(block)
+/**
+ * Generates lazy sequence.
+ */
+public fun <T> generate(block: suspend Generator<T>.() -> Unit): Sequence<T> = object : Sequence<T> {
+    override fun iterator(): Iterator<T> {
+        val iterator = GeneratorIterator<T>()
+        iterator.nextStep = block.createCoroutine(receiver = iterator, completion = iterator)
+        return iterator
+    }
 }
 
-private class GeneratedIterator<T>(block: suspend Generator<T>.() -> Unit) : AbstractIterator<T>(), Generator<T> {
-    private var nextStep: Continuation<Unit> = block.createCoroutine(this, object : Continuation<Unit> {
-        override fun resume(data: Unit) {
-            done()
+private class GeneratorIterator<T>: Generator<T>(), Iterator<T>, Continuation<Unit> {
+    var computedNext = false
+    var nextStep: Continuation<Unit>? = null
+    var nextValue: T? = null
+
+    override fun hasNext(): Boolean {
+        if (!computedNext) {
+            val step = nextStep!!
+            computedNext = true
+            nextStep = null
+            step.resume(Unit) // leaves it in "done" state if crashes
         }
+        return nextStep != null
+    }
 
-        override fun resumeWithException(exception: Throwable) {
-            throw exception
+    override fun next(): T {
+        if (!hasNext()) throw NoSuchElementException()
+        computedNext = false
+        return nextValue as T
+    }
+
+    // Completion continuation implementation
+    override fun resume(value: Unit) {
+        // nothing to do here -- leave null in nextStep
+    }
+
+    override fun resumeWithException(exception: Throwable) {
+        throw exception // just rethrow
+    }
+
+    // Generator implementation
+    override suspend fun yield(value: T) {
+        nextValue = value
+        return CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+            nextStep = c
+            CoroutineIntrinsics.SUSPENDED
         }
-    })
+    }
 
-    override fun computeNext() {
-        nextStep.resume(Unit)
+    override suspend fun yieldAll(iterator: Iterator<T>) {
+        if (!iterator.hasNext()) return
+        nextValue = iterator.next()
+        return CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+            nextStep = IteratorContinuation(c, iterator)
+            CoroutineIntrinsics.SUSPENDED
+        }
     }
-    suspend override fun yield(value: T) = CoroutineIntrinsics.suspendCoroutineOrReturn <Unit> { c ->
-        setNext(value)
-        nextStep = c
 
-        SUSPENDED
+    inner class IteratorContinuation(val completion: Continuation<Unit>, val iterator: Iterator<T>) : Continuation<Unit> {
+        override fun resume(value: Unit) {
+            if (!iterator.hasNext()) {
+                completion.resume(Unit)
+                return
+            }
+            nextValue = iterator.next()
+            nextStep = this
+        }
+
+        override fun resumeWithException(exception: Throwable) {
+            throw exception // just rethrow
+        }
     }
-}
+}

kotlinx-coroutines-generate/src/test/kotlin/GenerateTest.kt

@@ -147,4 +147,128 @@ class GenerateTest {
 
         assertEquals(listOf(Pair(1, 2), Pair(6, 8), Pair(15, 18)), result.zip(result).toList())
     }
+
+    @Test
+    fun testYieldAllIterator() {
+        val result = generate {
+            yieldAll(listOf(1, 2, 3).iterator())
+        }
+        assertEquals(listOf(1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllSequence() {
+        val result = generate {
+            yieldAll(sequenceOf(1, 2, 3))
+        }
+        assertEquals(listOf(1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollection() {
+        val result = generate {
+            yieldAll(listOf(1, 2, 3))
+        }
+        assertEquals(listOf(1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedFirst() {
+        val result = generate {
+            yield(0)
+            yieldAll(listOf(1, 2, 3))
+        }
+        assertEquals(listOf(0, 1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedLast() {
+        val result = generate {
+            yieldAll(listOf(1, 2, 3))
+            yield(4)
+        }
+        assertEquals(listOf(1, 2, 3, 4), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedBoth() {
+        val result = generate {
+            yield(0)
+            yieldAll(listOf(1, 2, 3))
+            yield(4)
+        }
+        assertEquals(listOf(0, 1, 2, 3, 4), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedLong() {
+        val result = generate {
+            yield(0)
+            yieldAll(listOf(1, 2, 3))
+            yield(4)
+            yield(5)
+            yieldAll(listOf(6))
+            yield(7)
+            yieldAll(listOf())
+            yield(8)
+        }
+        assertEquals(listOf(0, 1, 2, 3, 4, 5, 6, 7, 8), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionOneEmpty() {
+        val result = generate<Int> {
+            yieldAll(listOf())
+        }
+        assertEquals(listOf(), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionManyEmpty() {
+        val result = generate<Int> {
+            yieldAll(listOf())
+            yieldAll(listOf())
+            yieldAll(listOf())
+        }
+        assertEquals(listOf(), result.toList())
+    }
+
+    @Test
+    fun testYieldAllSideEffects() {
+        val effects = arrayListOf<Any>()
+        val result = generate {
+            effects.add("a")
+            yieldAll(listOf(1, 2))
+            effects.add("b")
+            yieldAll(listOf())
+            effects.add("c")
+            yieldAll(listOf(3))
+            effects.add("d")
+            yield(4)
+            effects.add("e")
+            yieldAll(listOf())
+            effects.add("f")
+            yield(5)
+        }
+
+        for (res in result) {
+            effects.add("(") // marks step start
+            effects.add(res)
+            effects.add(")") // marks step end
+        }
+        assertEquals(
+                listOf(
+                        "a",
+                        "(", 1, ")",
+                        "(", 2, ")",
+                        "b", "c",
+                        "(", 3, ")",
+                        "d",
+                        "(", 4, ")",
+                        "e", "f",
+                        "(", 5, ")"
+                ),
+                effects.toList()
+        )
+    }
 }