From 560c184fd43a63ca2b557a0752a590e2394a9763 Mon Sep 17 00:00:00 2001
From: Ilya Gorbunov <ilya.gorbunov@jetbrains.com>
Date: Fri, 13 Jan 2017 05:58:48 +0300
Subject: [PATCH] Merge AbstractIterator and SequenceBuilder, remove hasNext
 call from next.

---
 .../src/kotlin/collections/SequenceBuilder.kt | 125 ++++++------------
 1 file changed, 39 insertions(+), 86 deletions(-)

diff --git a/libraries/stdlib/src/kotlin/collections/SequenceBuilder.kt b/libraries/stdlib/src/kotlin/collections/SequenceBuilder.kt
index ff9ff20c672..9a20112da9e 100644
--- a/libraries/stdlib/src/kotlin/collections/SequenceBuilder.kt
+++ b/libraries/stdlib/src/kotlin/collections/SequenceBuilder.kt
@@ -47,7 +47,7 @@ public abstract class SequenceBuilder<in T> internal constructor() {
 
 internal fun <T> buildIteratorImpl(builderAction: suspend SequenceBuilder<T>.() -> Unit): Iterator<T> {
     val iterator = YieldingIterator<T>()
-    iterator.nextStep = builderAction.createCoroutine(receiver = iterator.builder, completion = iterator)
+    iterator.nextStep = builderAction.createCoroutine(receiver = iterator, completion = iterator)
     return iterator
 }
 
@@ -55,49 +55,51 @@ internal fun <T> buildIteratorImpl(builderAction: suspend SequenceBuilder<T>.()
 
 
 private typealias State = Int
-private const val State_Failed: State = 0
-private const val State_Done: State = 1
+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_ManyReady: State = 3
-private const val State_NotReady: State = 4
+private const val State_Done: State = 3
+private const val State_Failed: 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 class YieldingIterator<T> : SequenceBuilder<T>(), Iterator<T>, Continuation<Unit> {
     private var state = State_NotReady
-    private var nextValue: Any? = null
+    private var nextValue: T? = null
+    private var nextIterator: Iterator<T>? = null
+    var nextStep: Continuation<Unit>? = null
 
     override fun hasNext(): Boolean {
-        require(state != State_Failed)
         while (true) {
             when (state) {
-                State_Failed,
+                State_NotReady -> {}
+                State_ManyReady ->
+                    if (nextIterator!!.hasNext()) return true else nextIterator = null
                 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()
+            val step = nextStep!!
+            nextStep = null
+            step.resume(Unit)
         }
     }
 
-    override fun next(): T {
-        if (!hasNext()) throw NoSuchElementException()
+    override tailrec fun next(): T {
         when (state) {
-            State_Ready -> {
-                state = State_NotReady
-                return nextValue as T
-            }
+            State_NotReady -> if (!hasNext()) throw NoSuchElementException() else return next()
             State_ManyReady -> {
-                val iterator = nextValue as Iterator<T>
+                val iterator = nextIterator!!
                 return iterator.next()
             }
+            State_Ready -> {
+                state = State_NotReady
+                val result = nextValue as T
+                nextValue = null
+                return result
+            }
+            State_Done -> throw NoSuchElementException()
             else -> throw unexpectedState()
         }
     }
@@ -105,80 +107,31 @@ private abstract class AbstractIterator<T>: Iterator<T> {
     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 {
+    suspend override fun yield(value: T) {
         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
-            }
+        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
-            }
+    suspend override fun yieldAll(iterator: Iterator<T>) {
+        if (!iterator.hasNext()) return
+        nextIterator = iterator
+        state = State_ManyReady
+        return CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+            nextStep = c
+            CoroutineIntrinsics.SUSPENDED
         }
     }
 
     // Completion continuation implementation
     override fun resume(value: Unit) {
-        done()
+        state = State_Done
     }
 
     override fun resumeWithException(exception: Throwable) {
         throw exception // just rethrow
     }
-
 }