Introduce buildSequence and buildIterator functions based on coroutines.

2017-01-12 03:29:38 +03:00
parent 3225ec8935
commit 141ad43039
4 changed files with 468 additions and 0 deletions
@@ -3,6 +3,11 @@

 package kotlin.collections

+/**
+ * Builds an [Iterator] lazily yielding values one by one.
+ */
+public fun <T> buildIterator(builderAction: suspend SequenceBuilder<T>.() -> Unit): Iterator<T> = buildIteratorImpl(builderAction)
+
 /**
 * Creates an [Iterator] for an [Enumeration], allowing to use it in `for` loops.
 */
@@ -0,0 +1,184 @@
+@file:kotlin.jvm.JvmMultifileClass
+@file:kotlin.jvm.JvmName("SequencesKt")
+package kotlin.sequences
+
+import kotlin.coroutines.*
+
+/**
+ *  Builds a [Sequence] lazily yielding values one by one.
+ */
+public fun <T> buildSequence(builderAction: suspend SequenceBuilder<T>.() -> Unit): Sequence<T> = Sequence { buildIteratorImpl(builderAction) }
+
+
+/**
+ * Builder for a [Sequence] or an [Iterator], provides [yield] and [yieldAll] suspension functions.
+ */
+@RestrictsSuspension
+public abstract class SequenceBuilder<in T> internal constructor() {
+    /**
+     * Yields a value to the [Iterator] being built.
+     */
+    public abstract suspend fun yield(value: T)
+
+    /**
+     * Yields all values from the `iterator` to the [Iterator] being built.
+     *
+     * The sequence of values returned by the given iterator can be potentially infinite.
+     */
+    public abstract suspend fun yieldAll(iterator: Iterator<T>)
+
+    /**
+     * Yields a collections of values to the [Iterator] being built.
+     */
+    public suspend fun yieldAll(elements: Iterable<T>) {
+        if (elements is Collection && elements.isEmpty()) return
+        return yieldAll(elements.iterator())
+    }
+
+    /**
+     * Yields potentially infinite sequence of values  to the [Iterator] being built.
+     *
+     * The sequence can be potentially infinite.
+     */
+    public suspend fun yieldAll(sequence: Sequence<T>) = yieldAll(sequence.iterator())
+}
+
+
+
+internal fun <T> buildIteratorImpl(builderAction: suspend SequenceBuilder<T>.() -> Unit): Iterator<T> {
+    val iterator = YieldingIterator<T>()
+    iterator.nextStep = builderAction.createCoroutine(receiver = iterator.builder, completion = iterator)
+    return iterator
+}
+
+
+
+
+private typealias State = Int
+private const val State_Failed: State = 0
+private const val State_Done: State = 1
+private const val State_Ready: State = 2
+private const val State_ManyReady: State = 3
+private const val State_NotReady: State = 4
+
+/**
+ * A base class to simplify implementing iterators so that implementations only have to implement [computeNext]
+ * to implement the iterator, calling [done] when the iteration is complete.
+ */
+// TODO: Merge to kotlin.collections.AbstractIterator
+private abstract class AbstractIterator<T>: Iterator<T> {
+    private var state = State_NotReady
+    private var nextValue: Any? = null
+
+    override fun hasNext(): Boolean {
+        require(state != State_Failed)
+        while (true) {
+            when (state) {
+                State_Failed,
+                State_Done -> return false
+                State_Ready -> return true
+                State_ManyReady -> if ((nextValue as Iterator<*>).hasNext()) return true
+                State_NotReady -> {}
+                else -> throw unexpectedState()
+            }
+
+            state = State_Failed
+            computeNext()
+        }
+    }
+
+    override fun next(): T {
+        if (!hasNext()) throw NoSuchElementException()
+        when (state) {
+            State_Ready -> {
+                state = State_NotReady
+                return nextValue as T
+            }
+            State_ManyReady -> {
+                val iterator = nextValue as Iterator<T>
+                return iterator.next()
+            }
+            else -> throw unexpectedState()
+        }
+    }
+
+    private fun unexpectedState(): Throwable = IllegalStateException("Unexpected state of the iterator: $state")
+
+
+    /**
+     * Computes the next item in the iterator.
+     *
+     * This callback method should call one of these two methods:
+     *
+     * * [setNext] with the next value of the iteration
+     * * [done] to indicate there are no more elements
+     *
+     * Failure to call either method will result in the iteration terminating with a failed state
+     */
+    abstract protected fun computeNext(): Unit
+
+    /**
+     * Sets the next value in the iteration, called from the [computeNext] function
+     */
+    protected fun setNext(value: T): Unit {
+        nextValue = value
+        state = State_Ready
+    }
+
+    /**
+     * Sets the iterator to retrieve the next values from, called from the [computeNext] function
+     */
+    protected fun setNextMultiple(iterator: Iterator<T>): Unit {
+        nextValue = iterator
+        state = State_ManyReady
+    }
+
+    /**
+     * Sets the state to done so that the iteration terminates.
+     */
+    protected fun done() {
+        state = State_Done
+    }
+}
+
+
+
+
+private class YieldingIterator<T> : AbstractIterator<T>(), Continuation<Unit> {
+    var nextStep: Continuation<Unit>? = null
+
+    override fun computeNext() {
+        val step = nextStep!!
+        nextStep = null
+        step.resume(Unit)
+    }
+
+    val builder: SequenceBuilder<T> = object : SequenceBuilder<T>() {
+        suspend override fun yield(value: T) {
+            setNext(value)
+            return CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+                nextStep = c
+                CoroutineIntrinsics.SUSPENDED
+            }
+        }
+
+        suspend override fun yieldAll(iterator: Iterator<T>) {
+            if (!iterator.hasNext()) return
+            setNextMultiple(iterator)
+            return CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+                nextStep = c
+                CoroutineIntrinsics.SUSPENDED
+            }
+        }
+    }
+
+    // Completion continuation implementation
+    override fun resume(value: Unit) {
+        done()
+    }
+
+    override fun resumeWithException(exception: Throwable) {
+        throw exception // just rethrow
+    }
+
+}
@@ -0,0 +1,270 @@
+package test.collections
+
+import org.junit.Test
+import kotlin.test.*
+
+class SequenceBuilderTest {
+    @Test
+    fun testSimple() {
+        val result = buildSequence {
+            for (i in 1..3) {
+                yield(2 * i)
+            }
+        }
+
+        assertEquals(listOf(2, 4, 6), result.toList())
+        // Repeated calls also work
+        assertEquals(listOf(2, 4, 6), result.toList())
+    }
+
+    @Test
+    fun testCallHasNextSeveralTimes() {
+        val result = buildSequence {
+            yield(1)
+        }
+
+        val iterator = result.iterator()
+
+        assertTrue(iterator.hasNext())
+        assertTrue(iterator.hasNext())
+        assertTrue(iterator.hasNext())
+
+        assertEquals(1, iterator.next())
+
+        assertFalse(iterator.hasNext())
+        assertFalse(iterator.hasNext())
+        assertFalse(iterator.hasNext())
+
+        assertTrue(assertFails { iterator.next() } is NoSuchElementException)
+    }
+
+    @Test
+    fun testManualIteration() {
+        val result = buildSequence {
+            yield(1)
+            yield(2)
+            yield(3)
+        }
+
+        val iterator = result.iterator()
+
+        assertTrue(iterator.hasNext())
+        assertTrue(iterator.hasNext())
+        assertEquals(1, iterator.next())
+
+        assertTrue(iterator.hasNext())
+        assertTrue(iterator.hasNext())
+        assertEquals(2, iterator.next())
+
+        assertEquals(3, iterator.next())
+
+        assertFalse(iterator.hasNext())
+        assertFalse(iterator.hasNext())
+
+        assertTrue(assertFails { iterator.next() } is NoSuchElementException)
+
+        assertEquals(1, result.iterator().next())
+    }
+
+    @Test
+    fun testEmptySequence() {
+        val result = buildSequence<Int> {}
+        val iterator = result.iterator()
+
+        assertFalse(iterator.hasNext())
+        assertFalse(iterator.hasNext())
+
+        assertTrue(assertFails { iterator.next() } is NoSuchElementException)
+    }
+
+    @Test
+    fun testLaziness() {
+        var sharedVar = -2
+        val result = buildSequence {
+            while (true) {
+                when (sharedVar) {
+                    -1 -> return@buildSequence
+                    -2 -> error("Invalid state: -2")
+                    else -> yield(sharedVar)
+                }
+            }
+        }
+
+        val iterator = result.iterator()
+
+        sharedVar = 1
+        assertTrue(iterator.hasNext())
+        assertEquals(1, iterator.next())
+
+        sharedVar = 2
+        assertTrue(iterator.hasNext())
+        assertEquals(2, iterator.next())
+
+        sharedVar = 3
+        assertTrue(iterator.hasNext())
+        assertEquals(3, iterator.next())
+
+        sharedVar = -1
+        assertFalse(iterator.hasNext())
+        assertTrue(assertFails { iterator.next() } is NoSuchElementException)
+    }
+
+    @Test
+    fun testExceptionInCoroutine() {
+        var sharedVar = -2
+        val result = buildSequence {
+            while (true) {
+                when (sharedVar) {
+                    -1 -> return@buildSequence
+                    -2 -> error("Invalid state: -2")
+                    else -> yield(sharedVar)
+                }
+            }
+        }
+
+        val iterator = result.iterator()
+
+        sharedVar = 1
+        assertEquals(1, iterator.next())
+
+        sharedVar = -2
+        assertTrue(assertFails { iterator.hasNext() } is IllegalStateException)
+    }
+
+    @Test
+    fun testParallelIteration() {
+        var inc = 0
+        val result = buildSequence {
+            for (i in 1..3) {
+                inc++
+                yield(inc * i)
+            }
+        }
+
+        assertEquals(listOf(Pair(1, 2), Pair(6, 8), Pair(15, 18)), result.zip(result).toList())
+    }
+
+    @Test
+    fun testYieldAllIterator() {
+        val result = buildSequence {
+            yieldAll(listOf(1, 2, 3).iterator())
+        }
+        assertEquals(listOf(1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllSequence() {
+        val result = buildSequence {
+            yieldAll(sequenceOf(1, 2, 3))
+        }
+        assertEquals(listOf(1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollection() {
+        val result = buildSequence {
+            yieldAll(listOf(1, 2, 3))
+        }
+        assertEquals(listOf(1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedFirst() {
+        val result = buildSequence {
+            yield(0)
+            yieldAll(listOf(1, 2, 3))
+        }
+        assertEquals(listOf(0, 1, 2, 3), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedLast() {
+        val result = buildSequence {
+            yieldAll(listOf(1, 2, 3))
+            yield(4)
+        }
+        assertEquals(listOf(1, 2, 3, 4), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedBoth() {
+        val result = buildSequence {
+            yield(0)
+            yieldAll(listOf(1, 2, 3))
+            yield(4)
+        }
+        assertEquals(listOf(0, 1, 2, 3, 4), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionMixedLong() {
+        val result = buildSequence {
+            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 = buildSequence<Int> {
+            yieldAll(listOf())
+        }
+        assertEquals(listOf(), result.toList())
+    }
+
+    @Test
+    fun testYieldAllCollectionManyEmpty() {
+        val result = buildSequence<Int> {
+            yieldAll(listOf())
+            yieldAll(listOf())
+            yieldAll(listOf())
+        }
+        assertEquals(listOf(), result.toList())
+    }
+
+    @Test
+    fun testYieldAllSideEffects() {
+        val effects = arrayListOf<Any>()
+        val result = buildSequence {
+            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()
+        )
+    }
+}
@@ -1423,6 +1423,7 @@ public final class kotlin/collections/CollectionsKt {
 	public static synthetic fun binarySearch$default (Ljava/util/List;Ljava/lang/Object;Ljava/util/Comparator;IIILjava/lang/Object;)I
 	public static final fun binarySearchBy (Ljava/util/List;Ljava/lang/Comparable;IILkotlin/jvm/functions/Function1;)I
 	public static synthetic fun binarySearchBy$default (Ljava/util/List;Ljava/lang/Comparable;IILkotlin/jvm/functions/Function1;ILjava/lang/Object;)I
+	public static final fun buildIterator (Lkotlin/jvm/functions/Function2;)Ljava/util/Iterator;
 	public static final fun collectionSizeOrDefault (Ljava/lang/Iterable;I)I
 	public static final fun collectionSizeOrNull (Ljava/lang/Iterable;)Ljava/lang/Integer;
 	public static final fun contains (Ljava/lang/Iterable;Ljava/lang/Object;)Z
@@ -2061,6 +2062,13 @@ public abstract interface class kotlin/sequences/Sequence {
 	public abstract fun iterator ()Ljava/util/Iterator;
 }

+public abstract class kotlin/sequences/SequenceBuilder {
+	public abstract fun yield (Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public final fun yieldAll (Ljava/lang/Iterable;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public abstract fun yieldAll (Ljava/util/Iterator;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public final fun yieldAll (Lkotlin/sequences/Sequence;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+
 public final class kotlin/sequences/SequencesKt {
 	public static final fun all (Lkotlin/sequences/Sequence;Lkotlin/jvm/functions/Function1;)Z
 	public static final fun any (Lkotlin/sequences/Sequence;)Z
@@ -2079,6 +2087,7 @@ public final class kotlin/sequences/SequencesKt {
 	public static final fun averageOfInt (Lkotlin/sequences/Sequence;)D
 	public static final fun averageOfLong (Lkotlin/sequences/Sequence;)D
 	public static final fun averageOfShort (Lkotlin/sequences/Sequence;)D
+	public static final fun buildSequence (Lkotlin/jvm/functions/Function2;)Lkotlin/sequences/Sequence;
 	public static final fun constrainOnce (Lkotlin/sequences/Sequence;)Lkotlin/sequences/Sequence;
 	public static final fun contains (Lkotlin/sequences/Sequence;Ljava/lang/Object;)Z
 	public static final fun count (Lkotlin/sequences/Sequence;)I