Taking sequence elements fulfilling a predicate then continuing from there in Kotlin - 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.

Answer 2

I believe there is a way to accomplish what you want using the standard library. Limit the sequence first and then groupBy the number of digits.
val Int.numberOfDigits
get() = this.toString().length
sequenceOf(1,22,333).takeWhile{ it.numberOfDigits < 3 }.groupBy{ it.numberOfDigits }.values
If you want to avoid the eager evaluation of groupBy you could use groupingBy instead and then reduce potentially leaving the accumulator blank.

Answer 3

ardenit's answer seems like the best reusable approach. Since taking "chunks" of a sequence requires some state it doesn't seem likely something easily done in a purely functional manner. Delegating the state to a separate class enveloping the sequence makes sense.
Here's a small snippet showing what I ended up using. This assumes the sequence will not be empty and is (technically) infinite or further results aren't requested at some point.
class ChunkedIterator<T>(seq: Sequence<T>) {
private val it = seq.iterator()
var next: T = it.next()
fun next(predicate: (T) -> Boolean): List<T> {
val result = ArrayList<T>();
while (predicate.invoke(next)) {
result.add(next)
next = it.next();
}
return result
}
}

Answer 4

one way you could achieve this is by getting an iterator from your your original sequence and then building a new sequence out of it for each "take" -
val itr = seq.iterator()
val batch1 = itr.asSequence().takeWhile { predicate1(it) }.toList()
val batch2 = itr.asSequence().takeWhile { predicate2(it) }.toList()