Allow dropping and taking a lot of elements in very long sequences

For example allow dropping and taking Int.MAX_VALUE elements from
a sequence that is a result of another Int.MAX_VALUE dropping.
This commit is contained in:
Ilya Gorbunov
2018-08-09 18:57:01 +03:00
parent eaa0902ea5
commit c44f62a3d4
3 changed files with 23 additions and 2 deletions
@@ -41,5 +41,21 @@ class IndexOverflowJVMTest {
@Test
fun dropTwiceMaxValue() {
val halfMax = (1 shl 30) + 1
val dropOnce = longCountSequence.drop(halfMax)
val dropTwice = dropOnce.drop(halfMax)
val expectedEnd = halfMax.toLong() * 2
assertEquals(expectedEnd, dropTwice.first())
val dropTake = dropOnce.take(halfMax + 1)
assertEquals(expectedEnd, dropTake.last())
}
}
@@ -421,8 +421,8 @@ internal class DropSequence<T>(
require(count >= 0) { "count must be non-negative, but was $count." }
}
override fun drop(n: Int): Sequence<T> = DropSequence(sequence, count + n)
override fun take(n: Int): Sequence<T> = SubSequence(sequence, count, count + n)
override fun drop(n: Int): Sequence<T> = (count + n).let { n1 -> if (n1 < 0) DropSequence(this, n) else DropSequence(sequence, n1) }
override fun take(n: Int): Sequence<T> = (count + n).let { n1 -> if (n1 < 0) TakeSequence(this, n) else SubSequence(sequence, count, n1) }
override fun iterator(): Iterator<T> = object : Iterator<T> {
val iterator = sequence.iterator()
@@ -152,6 +152,11 @@ public class SequenceTest {
assertEquals("13, 21, 34, 55, 89, 144, 233, 377, 610, 987, ...", fibonacci().drop(7).joinToString(limit = 10))
assertEquals("13, 21, 34, 55, 89, 144, 233, 377, 610, 987, ...", fibonacci().drop(3).drop(4).joinToString(limit = 10))
assertFailsWith<IllegalArgumentException> { fibonacci().drop(-1) }
val dropMax = fibonacci().drop(Int.MAX_VALUE)
val dropMore = dropMax.drop(Int.MAX_VALUE)
val takeMore = dropMax.take(Int.MAX_VALUE)
}
@Test fun take() {