From 698ffa2807cfd7c3e59f74724ed03d82674d9fa8 Mon Sep 17 00:00:00 2001
From: Ilya Gorbunov <ilya.gorbunov@jetbrains.com>
Date: Fri, 27 Mar 2015 17:42:05 +0300
Subject: [PATCH] A sequence constructed with initialValue can be iterated
 multiple times.

---
 .../stdlib/src/kotlin/collections/Sequence.kt | 48 +++++++++++++++++--
 .../stdlib/test/collections/SequenceTest.kt   |  4 +-
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/libraries/stdlib/src/kotlin/collections/Sequence.kt b/libraries/stdlib/src/kotlin/collections/Sequence.kt
index 1057443462c..b559f80147e 100644
--- a/libraries/stdlib/src/kotlin/collections/Sequence.kt
+++ b/libraries/stdlib/src/kotlin/collections/Sequence.kt
@@ -442,6 +442,7 @@ public class DropWhileSequence<T>(private val sequence: Sequence<T>,
  * A sequence which repeatedly calls the specified [producer] function and returns its return values, until
  * `null` is returned from [producer].
  */
+deprecated("Implementation detail. Use function sequence(nextFunction: () -> T?) instead.")
 public class FunctionSequence<T : Any>(private val producer: () -> T?) : Sequence<T> {
     override fun iterator(): Iterator<T> = object : Iterator<T> {
         var nextState: Int = -1 // -1 for unknown, 0 for done, 1 for continue
@@ -478,11 +479,42 @@ public class FunctionSequence<T : Any>(private val producer: () -> T?) : Sequenc
     }
 }
 
+private class GeneratorSequence<T: Any>(private val getInitialValue: () -> T?, private val getNextValue: (T) -> T?): Sequence<T> {
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
+        var nextItem: T? = getInitialValue()
+        var nextState: Int = if (nextItem == null) 0 else 1 // -1 for unknown, 0 for done, 1 for continue
+
+        private fun calcNext() {
+            nextItem = getNextValue(nextItem!!)
+            nextState = if (nextItem == null) 0 else 1
+        }
+
+        override fun next(): T {
+            if (nextState == -1)
+                calcNext()
+            if (nextState == 0)
+                throw NoSuchElementException()
+            val result = nextItem as T
+            // Clean next to avoid keeping reference on yielded instance
+            // need to keep state
+            // nextItem = null
+            nextState = -1
+            return result
+        }
+
+        override fun hasNext(): Boolean {
+            if (nextState == -1)
+                calcNext()
+            return nextState == 1
+        }
+    }
+}
+
 /**
  * Returns a sequence which invokes the function to calculate the next value on each iteration until the function returns `null`.
  */
 public fun <T : Any> sequence(nextFunction: () -> T?): Sequence<T> {
-    return FunctionSequence(nextFunction)
+    return GeneratorSequence(nextFunction, { nextFunction() })
 }
 
 deprecated("Use sequence() instead")
@@ -490,10 +522,20 @@ public fun <T : Any> stream(nextFunction: () -> T?): Sequence<T> = sequence(next
 
 /**
  * Returns a sequence which invokes the function to calculate the next value based on the previous one on each iteration
- * until the function returns `null`.
+ * until the function returns `null`. The sequence starts with the specified [initialValue].
+ *
+ * The sequence can be iterated multiple times, each time starting with the [initialValue].
  */
 public /*inline*/ fun <T : Any> sequence(initialValue: T, nextFunction: (T) -> T?): Sequence<T> =
-        sequence(nextFunction.toGenerator(initialValue))
+        GeneratorSequence({ initialValue }, nextFunction)
+
+/**
+ * Returns a sequence which invokes the function [initialValueFunction] to get the first item and then
+ * [nextFunction] to calculate the next value based on the previous one on each iteration
+ * until the function returns `null`. The sequence starts with the value returned by [initialValueFunction].
+ */
+public fun <T: Any> sequence(initialValueFunction: () -> T?, nextFunction: (T) -> T?): Sequence<T> =
+        GeneratorSequence(initialValueFunction, nextFunction)
 
 deprecated("Use sequence() instead")
 public /*inline*/ fun <T : Any> stream(initialValue: T, nextFunction: (T) -> T?): Sequence<T> = sequence(initialValue, nextFunction)
diff --git a/libraries/stdlib/test/collections/SequenceTest.kt b/libraries/stdlib/test/collections/SequenceTest.kt
index 1d7fd229035..f8265fa1a5c 100644
--- a/libraries/stdlib/test/collections/SequenceTest.kt
+++ b/libraries/stdlib/test/collections/SequenceTest.kt
@@ -176,7 +176,9 @@ public class SequenceTest {
 
     test fun streamFromFunctionWithInitialValue() {
         val values = sequence(3) { n -> if (n > 0) n - 1 else null }
-        assertEquals(listOf(3, 2, 1, 0), values.toList())
+        val expected = listOf(3, 2, 1, 0)
+        assertEquals(expected, values.toList())
+        assertEquals(expected, values.toList(), "Iterating sequence second time yields the same result")
     }
 
     private fun <T, C : MutableCollection<in T>> Sequence<T>.takeWhileTo(result: C, predicate: (T) -> Boolean): C {