From 24edbc36df6efc45876b4f95456e4142ffa855c2 Mon Sep 17 00:00:00 2001 From: Ilya Gorbunov Date: Fri, 12 Feb 2016 20:37:16 +0300 Subject: [PATCH] Create at most one wrapper sequence for adjacent drop/take operations on sequences. Based on PR #814 by Sargun Vohra #KT-10821 Fixed --- libraries/stdlib/src/generated/_Sequences.kt | 12 ++- .../src/kotlin/collections/Sequences.kt | 84 +++++++++++++++++-- .../stdlib/test/collections/SequenceTest.kt | 28 ++++++- .../src/templates/Filtering.kt | 12 ++- 4 files changed, 122 insertions(+), 14 deletions(-) diff --git a/libraries/stdlib/src/generated/_Sequences.kt b/libraries/stdlib/src/generated/_Sequences.kt index 6f048e9f4fc..59e3df6666b 100644 --- a/libraries/stdlib/src/generated/_Sequences.kt +++ b/libraries/stdlib/src/generated/_Sequences.kt @@ -289,7 +289,11 @@ public inline fun Sequence.singleOrNull(predicate: (T) -> Boolean): T? { */ public fun Sequence.drop(n: Int): Sequence { require(n >= 0, { "Requested element count $n is less than zero." }) - return if (n == 0) this else DropSequence(this, n) + return when { + n == 0 -> this + this is DropTakeSequence -> this.drop(n) + else -> DropSequence(this, n) + } } /** @@ -367,7 +371,11 @@ public inline fun > Sequence.filterTo(destinat */ public fun Sequence.take(n: Int): Sequence { require(n >= 0, { "Requested element count $n is less than zero." }) - return if (n == 0) emptySequence() else TakeSequence(this, n) + return when { + n == 0 -> emptySequence() + this is DropTakeSequence -> this.take(n) + else -> TakeSequence(this, n) + } } /** diff --git a/libraries/stdlib/src/kotlin/collections/Sequences.kt b/libraries/stdlib/src/kotlin/collections/Sequences.kt index 49b7696cb64..372ae26d7db 100644 --- a/libraries/stdlib/src/kotlin/collections/Sequences.kt +++ b/libraries/stdlib/src/kotlin/collections/Sequences.kt @@ -36,8 +36,10 @@ public fun sequenceOf(vararg elements: T): Sequence = if (elements.isEmpt */ public fun emptySequence(): Sequence = EmptySequence -private object EmptySequence : Sequence { +private object EmptySequence : Sequence, DropTakeSequence { override fun iterator(): Iterator = EmptyIterator + override fun drop(n: Int) = EmptySequence + override fun take(n: Int) = EmptySequence } /** @@ -247,18 +249,79 @@ internal class FlatteningSequence } } +/** + * A sequence that supports drop(n) and take(n) operations + */ +internal interface DropTakeSequence : Sequence { + fun drop(n: Int): Sequence + fun take(n: Int): Sequence +} + +/** + * A sequence that skips [startIndex] values from the underlying [sequence] + * and stops returning values right before [endIndex], i.e. stops at `endIndex - 1` + */ +internal class SubSequence ( + private val sequence: Sequence, + private val startIndex: Int, + private val endIndex: Int +): Sequence, DropTakeSequence { + + init { + require(startIndex >= 0) { "startIndex should be non-negative, but is $startIndex" } + require(endIndex >= 0) { "endIndex should be non-negative, but is $endIndex" } + require(endIndex >= startIndex) { "endIndex should be not less than startIndex, but was $endIndex < $startIndex"} + } + + private val count: Int get() = endIndex - startIndex + + override fun drop(n: Int): Sequence = if (n >= count) emptySequence() else SubSequence(sequence, startIndex + n, endIndex) + override fun take(n: Int): Sequence = if (n >= count) this else SubSequence(sequence, startIndex, startIndex + n) + + override fun iterator() = object : Iterator { + + val iterator = sequence.iterator() + var position = 0 + + // Shouldn't be called from constructor to avoid premature iteration + private fun drop() { + while(position < startIndex && iterator.hasNext()) { + iterator.next() + position++ + } + } + + override fun hasNext(): Boolean { + drop() + return (position < endIndex) && iterator.hasNext() + } + + override fun next(): T { + drop() + if (position >= endIndex) + throw NoSuchElementException() + position++ + return iterator.next() + } + } +} + /** * A sequence that returns at most [count] values from the underlying [sequence], and stops returning values * as soon as that count is reached. */ -internal class TakeSequence - constructor(private val sequence: Sequence, - private val count: Int - ) : Sequence { +internal class TakeSequence ( + private val sequence: Sequence, + private val count: Int +) : Sequence, DropTakeSequence { + init { require (count >= 0) { throw IllegalArgumentException("count should be non-negative, but is $count") } } + override fun drop(n: Int): Sequence = if (n >= count) emptySequence() else SubSequence(sequence, n, count) + override fun take(n: Int): Sequence = if (n >= count) this else TakeSequence(sequence, n) + override fun iterator(): Iterator = object : Iterator { var left = count val iterator = sequence.iterator(); @@ -326,14 +389,17 @@ internal class TakeWhileSequence * A sequence that skips the specified number of values from the underlying [sequence] and returns * all values after that. */ -internal class DropSequence - constructor(private val sequence: Sequence, - private val count: Int - ) : Sequence { +internal class DropSequence ( + private val sequence: Sequence, + private val count: Int +) : Sequence, DropTakeSequence { init { require (count >= 0) { throw IllegalArgumentException("count should be non-negative, but is $count") } } + override fun drop(n: Int): Sequence = DropSequence(sequence, count + n) + override fun take(n: Int): Sequence = SubSequence(sequence, count, count + n) + override fun iterator(): Iterator = object : Iterator { val iterator = sequence.iterator(); var left = count diff --git a/libraries/stdlib/test/collections/SequenceTest.kt b/libraries/stdlib/test/collections/SequenceTest.kt index 4124d4da86c..fd5663a30ac 100644 --- a/libraries/stdlib/test/collections/SequenceTest.kt +++ b/libraries/stdlib/test/collections/SequenceTest.kt @@ -111,13 +111,39 @@ public class SequenceTest { } @test fun drop() { + assertEquals(emptyList(), emptySequence().drop(1).toList()) + listOf(2, 3, 4, 5).let { assertEquals(it, it.asSequence().drop(0).toList()) } 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)) + assertTrue(assertFails { fibonacci().drop(-1) } is IllegalArgumentException) } @test fun take() { + assertEquals(emptyList(), emptySequence().take(1).toList()) + assertEquals(emptyList(), fibonacci().take(0).toList()) + assertEquals("0, 1, 1, 2, 3, 5, 8", fibonacci().take(7).joinToString()) - assertEquals("2, 3, 5, 8", fibonacci().drop(3).take(4).joinToString()) + assertEquals("0, 1, 1, 2", fibonacci().take(7).take(4).joinToString()) + assertEquals("0, 1, 1, 2", fibonacci().take(4).take(5).joinToString()) + + assertEquals(emptyList(), fibonacci().take(1).drop(1).toList()) + assertEquals(emptyList(), fibonacci().take(1).drop(2).toList()) + + assertTrue(assertFails { fibonacci().take(-1) } is IllegalArgumentException) + } + + @test fun subSequence() { + assertEquals(listOf(2, 3, 5, 8), fibonacci().drop(3).take(4).toList()) + assertEquals(listOf(2, 3, 5, 8), fibonacci().take(7).drop(3).toList()) + + val seq = fibonacci().drop(3).take(4) + + assertEquals(listOf(2, 3, 5, 8), seq.take(5).toList()) + assertEquals(listOf(2, 3, 5), seq.take(3).toList()) + + assertEquals(emptyList(), seq.drop(5).toList()) + assertEquals(listOf(8), seq.drop(3).toList()) + } @test fun dropWhile() { diff --git a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt index 8aa006d9114..25e5e57670d 100644 --- a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt +++ b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt @@ -55,7 +55,11 @@ fun filtering(): List { body(Sequences) { """ require(n >= 0, { "Requested element count $n is less than zero." }) - return if (n == 0) this else DropSequence(this, n) + return when { + n == 0 -> this + this is DropTakeSequence -> this.drop(n) + else -> DropSequence(this, n) + } """ } @@ -121,7 +125,11 @@ fun filtering(): List { body(Sequences) { """ require(n >= 0, { "Requested element count $n is less than zero." }) - return if (n == 0) emptySequence() else TakeSequence(this, n) + return when { + n == 0 -> emptySequence() + this is DropTakeSequence -> this.take(n) + else -> TakeSequence(this, n) + } """ }