diff --git a/libraries/stdlib/src/kotlin/Iterators.kt b/libraries/stdlib/src/kotlin/Iterators.kt index e501eeefb0d..ce64ba62e39 100644 --- a/libraries/stdlib/src/kotlin/Iterators.kt +++ b/libraries/stdlib/src/kotlin/Iterators.kt @@ -2,6 +2,7 @@ package kotlin import kotlin.support.* import java.util.Collections +import kotlin.test.assertTrue /** * Returns an iterator which invokes the function to calculate the next value on each iteration until the function returns *null* @@ -10,6 +11,23 @@ public inline fun iterate(nextFunction: () -> T?) : Iterator { return FunctionIterator(nextFunction) } +/** + * Returns an iterator which invokes the function to calculate the next value based on the previous one on each iteration + * until the function returns *null* + */ +public inline fun iterate(initialValue: T, nextFunction: (T) -> T?): Iterator = + iterate(nextFunction.toGenerator(initialValue)) + +/** + * Returns an iterator whose values are pairs composed of values produced by given pair of iterators + */ +public inline fun Iterator.zip(iterator: Iterator): Iterator> = PairIterator(this, iterator) + +/** + * Returns an iterator shifted to right by the given number of elements + */ +public inline fun Iterator.skip(n: Int): Iterator = SkippingIterator(this, n) + class FilterIterator(val iterator : Iterator, val predicate: (T)-> Boolean) : AbstractIterator() { override protected fun computeNext(): Unit { while (iterator.hasNext()) { @@ -147,3 +165,50 @@ class IndexIterator(val iterator : Iterator): Iterator> { return iterator.hasNext() } } + +public class PairIterator( + val iterator1 : Iterator, val iterator2 : Iterator +): AbstractIterator>() { + protected override fun computeNext() { + if (iterator1.hasNext() && iterator2.hasNext()) { + setNext(Pair(iterator1.next(), iterator2.next())) + } + else { + done() + } + } +} + +class SkippingIterator(val iterator: Iterator, val n: Int): Iterator { + private var firstTime: Boolean = true + + private fun skip() { + for (i in 1..n) { + if (!iterator.hasNext()) break + iterator.next() + } + firstTime = false + } + + override fun next(): T { + assertTrue(!firstTime, "hasNext() must be invoked before advancing an iterator") + return iterator.next() + } + + override fun hasNext(): Boolean { + if (firstTime) { + skip() + } + return iterator.hasNext() + } +} + +fun Function1.toGenerator(initialValue: T): Function0 { + var nextValue: T? = initialValue + return { + nextValue?.let { result -> + nextValue = this@toGenerator(result) + result + } + } +} \ No newline at end of file diff --git a/libraries/stdlib/test/iterators/FunctionIteratorTest.kt b/libraries/stdlib/test/iterators/FunctionIteratorTest.kt index b1aa318ff55..8b788be839b 100644 --- a/libraries/stdlib/test/iterators/FunctionIteratorTest.kt +++ b/libraries/stdlib/test/iterators/FunctionIteratorTest.kt @@ -18,4 +18,9 @@ class FunctionIteratorTest { val list = iter.toList() assertEquals(arrayList(2, 1, 0), list) } + + Test fun iterateOverFunction2() { + val values = iterate(3) { n -> if (n > 0) n - 1 else null } + assertEquals(arrayList(3, 2, 1, 0), values.toList()) + } } \ No newline at end of file diff --git a/libraries/stdlib/test/iterators/IteratorsTest.kt b/libraries/stdlib/test/iterators/IteratorsTest.kt index fae3fc23b97..831678ea997 100644 --- a/libraries/stdlib/test/iterators/IteratorsTest.kt +++ b/libraries/stdlib/test/iterators/IteratorsTest.kt @@ -77,4 +77,14 @@ class IteratorsTest { assertEquals("13, 21, 34, 55, 89, 144, 233, 377, 610, 987, ...", fibonacci().filter { it > 10 }.makeString(limit = 10)) assertEquals("144, 233, 377, 610, 987", fibonacci().filter { it > 100 }.takeWhile { it < 1000 }.makeString()) } + + test fun pairIterator() { + val pairStr = (fibonacci() zip fibonacci().map { i -> i*2 }).makeString(limit = 10) + assertEquals("(0, 0), (1, 2), (1, 2), (2, 4), (3, 6), (5, 10), (8, 16), (13, 26), (21, 42), (34, 68), ...", pairStr) + } + + test fun skippingIterator() { + assertEquals("13, 21, 34, 55, 89, 144, 233, 377, 610, 987, ...", fibonacci().skip(7).makeString(limit = 10)) + assertEquals("13, 21, 34, 55, 89, 144, 233, 377, 610, 987, ...", fibonacci().skip(3).skip(4).makeString(limit = 10)) + } }