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

@@ -289,7 +289,11 @@ public inline fun <T> Sequence<T>.singleOrNull(predicate: (T) -> Boolean): T? {
  */
 public fun <T> Sequence<T>.drop(n: Int): Sequence<T> {
     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 <T, C : MutableCollection<in T>> Sequence<T>.filterTo(destinat
  */
 public fun <T> Sequence<T>.take(n: Int): Sequence<T> {
     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)
+    }
 }
 
 /**

libraries/stdlib/src/kotlin/collections/Sequences.kt

@@ -36,8 +36,10 @@ public fun <T> sequenceOf(vararg elements: T): Sequence<T> = if (elements.isEmpt
  */
 public fun <T> emptySequence(): Sequence<T> = EmptySequence
 
-private object EmptySequence : Sequence<Nothing> {
+private object EmptySequence : Sequence<Nothing>, DropTakeSequence<Nothing> {
     override fun iterator(): Iterator<Nothing> = EmptyIterator
+    override fun drop(n: Int) = EmptySequence
+    override fun take(n: Int) = EmptySequence
 }
 
 /**
@@ -247,18 +249,79 @@ internal class FlatteningSequence<T, R, E>
     }
 }
 
+/**
+ * A sequence that supports drop(n) and take(n) operations
+ */
+internal interface DropTakeSequence<T> : Sequence<T> {
+    fun drop(n: Int): Sequence<T>
+    fun take(n: Int): Sequence<T>
+}
+
+/**
+ * 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<T> (
+        private val sequence: Sequence<T>,
+        private val startIndex: Int,
+        private val endIndex: Int
+): Sequence<T>, DropTakeSequence<T> {
+
+    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<T> = if (n >= count) emptySequence() else SubSequence(sequence, startIndex + n, endIndex)
+    override fun take(n: Int): Sequence<T> = if (n >= count) this else SubSequence(sequence, startIndex, startIndex + n)
+
+    override fun iterator() = object : Iterator<T> {
+
+        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<T>
-                            constructor(private val sequence: Sequence<T>,
-                             private val count: Int
-                            ) : Sequence<T> {
+internal class TakeSequence<T> (
+    private val sequence: Sequence<T>,
+    private val count: Int
+) : Sequence<T>, DropTakeSequence<T> {
+
     init {
         require (count >= 0) { throw IllegalArgumentException("count should be non-negative, but is $count") }
     }
 
+    override fun drop(n: Int): Sequence<T> = if (n >= count) emptySequence() else SubSequence(sequence, n, count)
+    override fun take(n: Int): Sequence<T> = if (n >= count) this else TakeSequence(sequence, n)
+
     override fun iterator(): Iterator<T> = object : Iterator<T> {
         var left = count
         val iterator = sequence.iterator();
@@ -326,14 +389,17 @@ internal class TakeWhileSequence<T>
  * A sequence that skips the specified number of values from the underlying [sequence] and returns
  * all values after that.
  */
-internal class DropSequence<T>
-                            constructor(private val sequence: Sequence<T>,
-                             private val count: Int
-                            ) : Sequence<T> {
+internal class DropSequence<T> (
+        private val sequence: Sequence<T>,
+        private val count: Int
+) : Sequence<T>, DropTakeSequence<T> {
     init {
         require (count >= 0) { throw IllegalArgumentException("count should be non-negative, but is $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 iterator(): Iterator<T> = object : Iterator<T> {
         val iterator = sequence.iterator();
         var left = count

libraries/stdlib/test/collections/SequenceTest.kt

@@ -111,13 +111,39 @@ public class SequenceTest {
     }
 
     @test fun drop() {
+        assertEquals(emptyList(), emptySequence<Int>().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<Int>().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() {

libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt

@@ -55,7 +55,11 @@ fun filtering(): List<GenericFunction> {
         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<GenericFunction> {
         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)
+            }
             """
         }