diff --git a/libraries/stdlib/src/kotlin/collections/Sequence.kt b/libraries/stdlib/src/kotlin/collections/Sequence.kt index a7fb281eca1..3727793cfac 100644 --- a/libraries/stdlib/src/kotlin/collections/Sequence.kt +++ b/libraries/stdlib/src/kotlin/collections/Sequence.kt @@ -473,29 +473,28 @@ private class DistinctIterator(private val source : Iterator, private v private class GeneratorSequence(private val getInitialValue: () -> T?, private val getNextValue: (T) -> T?): Sequence { override fun iterator(): Iterator = object : Iterator { - var nextItem: T? = getInitialValue() - var nextState: Int = if (nextItem == null) 0 else 1 // -1 for unknown, 0 for done, 1 for continue + var nextItem: T? = null + var nextState: Int = -2 // -2 for initial unknown, -1 for next unknown, 0 for done, 1 for continue private fun calcNext() { - nextItem = getNextValue(nextItem!!) + nextItem = if (nextState == -2) getInitialValue() else getNextValue(nextItem!!) nextState = if (nextItem == null) 0 else 1 } override fun next(): T { - if (nextState == -1) + if (nextState < 0) 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 + // Do not clean nextItem (to avoid keeping reference on yielded instance) -- need to keep state for getNextValue nextState = -1 return result } override fun hasNext(): Boolean { - if (nextState == -1) + if (nextState < 0) calcNext() return nextState == 1 } @@ -530,7 +529,10 @@ public fun sequence(nextFunction: () -> T?): Sequence { * * The sequence can be iterated multiple times, each time starting with the [initialValue]. */ -public /*inline*/ fun sequence(initialValue: T, nextFunction: (T) -> T?): Sequence = +public fun sequence(initialValue: T?, nextFunction: (T) -> T?): Sequence = + if (initialValue == null) + EmptySequence + else GeneratorSequence({ initialValue }, nextFunction) /** diff --git a/libraries/stdlib/test/collections/SequenceTest.kt b/libraries/stdlib/test/collections/SequenceTest.kt index 559301263dd..4b8382cf958 100644 --- a/libraries/stdlib/test/collections/SequenceTest.kt +++ b/libraries/stdlib/test/collections/SequenceTest.kt @@ -1,6 +1,7 @@ package test.collections import org.junit.Test as test +import org.junit.Test import kotlin.test.* import java.util.* @@ -233,6 +234,24 @@ public class SequenceTest { assertEquals(expected, values.toList(), "Iterating sequence second time yields the same result") } + @Test fun sequenceFromFunctionWithLazyInitialValue() { + var start = 3 + val values = sequence({ start }, { n -> if (n > 0) n - 1 else null }) + val expected = listOf(3, 2, 1, 0) + assertEquals(expected, values.toList()) + assertEquals(expected, values.toList(), "Iterating sequence second time yields the same result") + + start = 2 + assertEquals(expected.drop(1), values.toList(), "Initial value function is called on each iterator request") + + // does not throw on construction + val errorValues = sequence({ (throw IllegalStateException()) }, { null }) + // does not throw on iteration + val iterator = errorValues.iterator() + // throws on advancing + assertFails { iterator.next() } + } + @test fun sequenceFromIterator() { val list = listOf(3, 2, 1, 0)