From c44f62a3d407c9890d72e312c57c85c158762f93 Mon Sep 17 00:00:00 2001
From: Ilya Gorbunov <ilya.gorbunov@jetbrains.com>
Date: Thu, 9 Aug 2018 18:57:01 +0300
Subject: [PATCH] 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.
---
 .../jvm/test/collections/IndexOverflowJVMTest.kt | 16 ++++++++++++++++
 .../stdlib/src/kotlin/collections/Sequences.kt   |  4 ++--
 .../stdlib/test/collections/SequenceTest.kt      |  5 +++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/libraries/stdlib/jvm/test/collections/IndexOverflowJVMTest.kt b/libraries/stdlib/jvm/test/collections/IndexOverflowJVMTest.kt
index c6f7b56e109..94807b9826d 100644
--- a/libraries/stdlib/jvm/test/collections/IndexOverflowJVMTest.kt
+++ b/libraries/stdlib/jvm/test/collections/IndexOverflowJVMTest.kt
@@ -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())
+    }
 
 }
\ No newline at end of file
diff --git a/libraries/stdlib/src/kotlin/collections/Sequences.kt b/libraries/stdlib/src/kotlin/collections/Sequences.kt
index 5f617fa99d1..f23d0d64bfa 100644
--- a/libraries/stdlib/src/kotlin/collections/Sequences.kt
+++ b/libraries/stdlib/src/kotlin/collections/Sequences.kt
@@ -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()
diff --git a/libraries/stdlib/test/collections/SequenceTest.kt b/libraries/stdlib/test/collections/SequenceTest.kt
index 01338c987c8..b740931ff46 100644
--- a/libraries/stdlib/test/collections/SequenceTest.kt
+++ b/libraries/stdlib/test/collections/SequenceTest.kt
@@ -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() {