How can I take varying chunks out of a Kotlin Sequence? - kotlin

Question

If I have a Kotlin sequence, every invocation of take(n) restarts the sequence.
val items = generateSequence(0) {
if (it > 9) null else it + 1
}
#Test fun `take doesn't remember position`() {
assertEquals(listOf(0, 1), items.take(2).toList())
assertEquals(listOf(0, 1, 2), items.take(3).toList())
}
Is there an easy way of write say, another(n) such that
#Test fun `another does remember position`() {
assertEquals(listOf(0, 1), items.another(2).toList())
assertEquals(listOf(2, 3, 4), items.another(3).toList())
}
I suppose that I have to have something that isn't the Sequence to keep the state, so maybe what I'm actually asking for is a nice definition of fun Iterator<T>.another(count: Int): List<T>

Answer 1

Sequence does not remember its position, but its iterator does remember:
val iterator : Iterator<Int> = items.iterator()
Now all you need is something like take(n) but for Iterator<T>:
public fun <T> Iterator<T>.another(n: Int): List<T> {
require(n >= 0) { "Requested element count $n is less than zero." }
if (n == 0) return emptyList()
var count = 0
val list = ArrayList<T>(n)
for (item in this) {
list.add(item)
if (++count == n)
break
}
return list
}

Answer 2

What about this:
#Test
fun `another does remember position`() {
val items: Sequence<Int> = generateSequence(0) {
if (it > 9) null else it + 1
}
val (first, rest) = items.another(2)
assertEquals(listOf(0, 1), first.toList())
assertEquals(listOf(2, 3, 4), rest.another(3).first.toList())
}
fun <T> Sequence<T>.another(n: Int): Pair<Sequence<T>, Sequence<T>> {
return this.take(n) to this.drop(n)
}

Answer 3

To answer the last part of your question:
I suppose that I have to have something that isn't the Sequence to keep the state, so maybe what I'm actually asking for is a nice definition of fun Iterator.another(count: Int): List
One such implementation would be:
fun <T> Iterator<T>.another(count: Int): List<T> {
val collectingList = mutableListOf<T>()
while (hasNext() && collectingList.size < count) {
collectingList.add(next())
}
return collectingList.toList()
}
This passes your test if you use the iterator produced by the sequence:
#Test
fun `another does remember position`() {
val items = generateSequence(0) {
if (it > 9) null else it + 1
}.iterator() //Use the iterator of this sequence.
assertEquals(listOf(0, 1), items.another(2))
assertEquals(listOf(2, 3, 4), items.another(3))
}
To me what you've described is an iterator, since it's something that allows you to go over a collection or sequence etc. but also remember its last position.
NB the implementation above wasn't written to take into consideration what should happen for non-positive counts passed in, and if the count is larger than what's left to iterate over you'll be returned a list which has smaller size than n. I suppose you could consider this an exercise for yourself :-)

Answer 4

Sequence does not remember its position, but its iterator does remember:
val iterator : Iterator<Int> = items.iterator()
Unfortunately there is no take(n) for an iterator, so to use the one from stdlib you need to wrap iter into an Iterable:
val iterable : Iterable<Int> = items.iterator().asIterable()
fun <T> Iterator<T>.asIterable() : Iterable<T> = object : Iterable<T> {
private val iter = this#asIterable
override fun iterator() = iter
}
That makes itareble.take(n) remember its position, but unfortunately there is a of-by-one error because the standard .take(n) asks for one element too many:
public fun <T> Iterable<T>.take(n: Int): List<T> {
require(n >= 0) { "Requested element count $n is less than zero." }
if (n == 0) return emptyList()
if (this is Collection<T>) {
if (n >= size) return toList()
if (n == 1) return listOf(first())
}
var count = 0
val list = ArrayList<T>(n)
for (item in this) {
if (count++ == n)
break
list.add(item)
}
return list.optimizeReadOnlyList()
}
That can be fixed with a little tweak:
public fun <T> Iterable<T>.take2(n: Int): List<T> {
require(n >= 0) { "Requested element count $n is less than zero." }
if (n == 0) return emptyList()
if (this is Collection<T>) {
if (n >= size) return toList()
if (n == 1) return listOf(first())
}
var count = 0
val list = ArrayList<T>(n)
for (item in this) {
list.add(item)
//count++
if (++count == n)
break
}
return list
}
Now both of you tests pass:
#Test fun `take does not remember position`() {
assertEquals(listOf(0, 1), items.take2(2).toList())
assertEquals(listOf(0, 1, 2), items.take2(3).toList())
}
#Test fun `another does remember position`() {
assertEquals(listOf(0, 1), iter.take2(2).toList())
assertEquals(listOf(2, 3, 4), iter.take2(3).toList())
}

Answer 5

You could create a function generateStatefulSequence which creates a sequence which keeps its state by using a second sequence's iterator to provide the values.
The iterator is captured in the closure of that function.
On each iteration the seed lambda ({ i.nextOrNull() }) of the returned sequence starts off with the next value provided by the iterator.
// helper
fun <T> Iterator<T>.nextOrNull() = if(hasNext()) { next() } else null
fun <T : Any> generateStatefulSequence(seed: T?, nextFunction: (T) -> T?): Sequence<T> {
val i = generateSequence(seed) {
nextFunction(it)
}.iterator()
return generateSequence(
seedFunction = { i.nextOrNull() },
nextFunction = { i.nextOrNull() }
)
}
Usage:
val s = generateStatefulSequence(0) { if (it > 9) null else it + 1 }
println(s.take(2).toList()) // [0, 1]
println(s.take(3).toList()) // [2, 3, 4]
println(s.take(10).toList()) // [5, 6, 7, 8, 9, 10]
Try it out

Answer 6

Here is a nice definition of fun Iterator<T>.another(count: Int): List<T> as requested:
fun <T> Iterator<T>.another(count: Int): List<T> =
if (count > 0 && hasNext()) listOf(next()) + this.another(count - 1)
else emptyList()

Answer 7

As another workaround (similar to the suggestion by Willi Mentzel above) would be to create a asStateful() extension method that converts any sequence into a one that will remember the position, by wrapping it into an Iterable that always yields the same iterator.
class StatefulIterable<out T>(wrapped: Sequence<T>): Iterable<T> {
private val iterator = wrapped.iterator()
override fun iterator() = iterator
}
fun <T> Sequence<T>.asStateful(): Sequence<T> = StatefulIterable(this).asSequence()
Then you can do:
val items = generateSequence(0) {
if (it > 9) null else it + 1
}.asStateful()
#Test fun `stateful sequence does remember position`() {
assertEquals(listOf(0, 1), items.take(2).toList())
assertEquals(listOf(2, 3, 4), items.take(3).toList())
}
Try it here: https://pl.kotl.in/Yine8p6wn