Taking sequence elements fulfilling a predicate then continuing from there in Kotlin

Question

In Kotlin sequences have a takeWhile function that will let you take items as long as they adhere to a given predicate. What I'd like to do is take items according to that predicate, use them in some way, then alter the predicate and take the next "batch". So far I haven't really found a way of doing this purely with what sequences and iterators offer.

Following snippet of code illustrates the problem. The primeGenerator() function returns a Sequence of prime (Long) numbers. Suppose that I want to make lists with each list having prime numbers with the same number of digits. On creating each list I'd use it for some purpose. If the list conforms to what I was searching the iteration can end, otherwise move onto the next list.

val primeIt = primeGenerator().iterator()
var digits = 1
var next: Long? = null
val currentList = ArrayList<Long>()
while (digits < 4) {
    next?.also { currentList.add(it) }
    next = primeIt.next()
    if (next.toString().length > digits) {
        println("Primes with $digits: $currentList")
        currentList.clear()
        digits++
    }
}

In this case it ends once the number of digits exceeds 3. This works fine, but I was wondering if there is some way to achieve the same with operations chained purely on the sequence or an iterator of it. Basically chunking the sequence but based on a predicate rather than a set size. The prime number example above is just for illustration, I'm after the general principle, not something that'd only work for this case.

Answer 1

There are no such functions in standard library for large (or infinite) sequences, but you may write such function by yourself (although it requires some extra code):

class BufferedIterator<T>(private val iterator: Iterator<T>) : Iterator<T> {

    var current: T? = null
        private set

    var reachedEnd: Boolean = false
        private set

    override fun hasNext(): Boolean = iterator.hasNext().also { reachedEnd = !it }

    override fun next(): T = iterator.next().also { current = it }
}

fun <T> Iterator<T>.buffered() = BufferedIterator(this)

fun <T> BufferedIterator<T>.takeWhile(predicate: (T) -> Boolean): List<T> {
    val list = ArrayList<T>()
    if (reachedEnd) return list
    current?.let {
        if (predicate(it)) list += it
    }
    while (hasNext()) {
        val next = next()
        if (predicate(next)) list += next
        else break
    }
    return list
}

fun main() {
    val sequence = sequence {
        var next = 0
        while (true) {
            yield(next++)
        }
    }
    val iter = sequence.iterator().buffered()
    for (i in 0..3) {
        println(iter.takeWhile { it.toString().length <= i })
    }
}

With this approach you can easily work even with infinite sequences.