diff --git a/libraries/stdlib/src/Iterators.kt b/libraries/stdlib/src/Iterators.kt index 99f82c0dfdf..8ff451c349a 100644 --- a/libraries/stdlib/src/Iterators.kt +++ b/libraries/stdlib/src/Iterators.kt @@ -1,6 +1,7 @@ package kotlin import kotlin.support.AbstractIterator +import kotlin.support.FunctionIterator import java.util.* import java.util.Iterator @@ -77,3 +78,8 @@ inline fun java.util.Iterator.flatMap(transform: (T)-> java.util.Itera return flatMapTo<>(ArrayList(), transform) } */ + +/** + * Returns an iterator which invokes the function to calculate the next value on each iteration until the function returns null + */ +inline fun iterate(nextFunction: () -> T?) : Iterator = FunctionIterator(nextFunction) \ No newline at end of file diff --git a/libraries/stdlib/src/support/AbstractIterator.kt b/libraries/stdlib/src/support/AbstractIterator.kt index 440301390d1..4c7146f1631 100644 --- a/libraries/stdlib/src/support/AbstractIterator.kt +++ b/libraries/stdlib/src/support/AbstractIterator.kt @@ -70,4 +70,19 @@ abstract class AbstractIterator: java.util.Iterator { protected fun done() { state = State.Done } +} + +/** + * An [[Iterator]] implementation which invokes a function to calculate the next value in the iteration + * until the function returns null + */ +class FunctionIterator(val nextFn: () -> T?) : AbstractIterator() { + + override fun computeNext(): T? { + val next = (nextFn)() + if (next == null) { + done() + } + return next + } } \ No newline at end of file diff --git a/libraries/stdlib/test/iterators/FunctionIteratorTest.kt b/libraries/stdlib/test/iterators/FunctionIteratorTest.kt new file mode 100644 index 00000000000..98ab172fab2 --- /dev/null +++ b/libraries/stdlib/test/iterators/FunctionIteratorTest.kt @@ -0,0 +1,22 @@ +package iterators + +import kotlin.* +import kotlin.test.* +import kotlin.util.* + +import org.junit.Test + +class FunctionIteratorTest { + + Test fun iterateOverFunction() { + var count = 3 + + val iter = iterate { + count-- + if (count >= 0) count else null + } + + val list = iter.toList() + assertEquals(arrayList(2, 1, 0), list) + } +} \ No newline at end of file