iterate.kt

/**
 * My take on the kotlin iterator presented in this (outdated) document:
 *
 * http://jamie.mccrindle.org/2013/01/a-functional-iterate-for-kotlin.html
 *
 */

import org.junit.Test
import kotlin.support.AbstractIterator
import kotlin.test.assertEquals
import kotlin.test.assertFalse

class iterateTest {
    @Test
    fun iterateWorks() {
        var i = 5
        var iterator: Iterator<Any> = iterate { if (i == 8) null else i++ }
        assertEquals(5, iterator.next())
        assertEquals(6, iterator.next())
        assertEquals(7, iterator.next())
        assertFalse(iterator.hasNext())

        iterator = listOf(1,2,3).iterator().map { "hello $it" }
        assertEquals("hello 1", iterator.next())
        assertEquals("hello 2", iterator.next())
        assertEquals("hello 3", iterator.next())

        iterator = iterate(0, { it + 1 })
        assertEquals(0, iterator.next())
        assertEquals(1, iterator.next())
        assertEquals(2, iterator.next())

        iterator = iterate(0, { it + 2 }).map { it + 0.5 }
        assertEquals(0.5, iterator.next())
        assertEquals(2.5, iterator.next())
        assertEquals(4.5, iterator.next())
        
        iterator = iterate(5, { it + 2 }).map { "I see $it cats" }
        assertEquals("I see 5 cats", iterator.next())
        assertEquals("I see 7 cats", iterator.next())
        assertEquals("I see 9 cats", iterator.next())
    }
}

/**
 * A functional implementation of Iterator. Pass a function that returns null when done
 * or the next value in the sequence.
 *
 * <pre class="prettyprint">
 *   var i = 0
 *   val iterator = iterate({ i++ })
 * </pre>
 */
fun <T> iterate(next: () -> T?): Iterator<T> {
    return object : AbstractIterator<T>() {
        override fun computeNext() {
            val value = next.invoke() as T
            if (value == null) {
                done()
            }
            else {
                setNext(value)
            }
        }
    }
}

/**
 * Returns an iterator. This method will keep calling next with the current state and then
 * updating the current state with the result, starting with the initial state. If next
 * returns null, close the iterator.
 *
 * <pre class="prettyprint">
 *   val iterator = iterate(5, { it + 1 })
 * </pre>
 */
fun <T> iterate(initial: T?, next: (T) -> T?): Iterator<T> {
    var current = initial
    return iterate {
        if (current != null) {
            val result = current as T
            current = next(result)
            result
        }
        else {
            null
        }
    }
}

/**
 * Returns a new iterator, in which every output value from the current iterator
 * will be passed through the mapping function.
 */
fun <T, R> Iterator<T>.map(f: (T) -> R): Iterator<R> {
    return iterate { f.invoke(next()) }
}