Make sure index and count do not overflow for long sequences

Throw an exception immediately before an overflow becomes observable.
Place check to prevent negative index from indexOf, indexOfFirst.
Do not insert overflow checks for arrays, lists, maps and char sequences.

#KT-16097

libraries/stdlib/jvm/src/kotlin/collections/CollectionsJVM.kt

@@ -9,6 +9,8 @@
 package kotlin.collections
 
 import kotlin.*
+import kotlin.internal.InlineOnly
+import kotlin.internal.apiVersionIsAtLeast
 
 /**
  * Returns an immutable list containing only the specified object [element].
@@ -42,3 +44,31 @@ internal actual fun <T> Array<out T>.copyToArrayOfAny(isVarargs: Boolean): Array
         @Suppress("UNCHECKED_CAST") (this as Array<Any?>)
     else
         java.util.Arrays.copyOf(this, this.size, Array<Any?>::class.java)
+
+
+@PublishedApi
+@SinceKotlin("1.3")
+@InlineOnly
+internal actual inline fun checkIndexOverflow(index: Int): Int {
+    if (index < 0) {
+        if (apiVersionIsAtLeast(1, 3, 0))
+            throwIndexOverflow()
+        else
+            throw ArithmeticException("Index overflow has happened.")
+    }
+    return index
+}
+
+@PublishedApi
+@SinceKotlin("1.3")
+@InlineOnly
+internal actual inline fun checkCountOverflow(count: Int): Int {
+    if (count < 0) {
+        if (apiVersionIsAtLeast(1, 3, 0))
+            throwCountOverflow()
+        else
+            throw ArithmeticException("Count overflow has happened.")
+    }
+    return count
+}
+

libraries/stdlib/jvm/test/collections/IndexOverflowJVMTest.kt

@@ -37,8 +37,114 @@ class IndexOverflowJVMTest {
                 override fun next(): Long = counter++
             }
         }
+
+        fun assertIndexOverflow(f: () -> Unit) {
+            val ex = assertFailsWith<ArithmeticException>(block = f)
+            assertTrue(ex.message!!.contains("index", ignoreCase = true))
+        }
+
+        fun assertCountOverflow(f: () -> Unit) {
+            val ex = assertFailsWith<ArithmeticException>(block = f)
+            assertTrue(ex.message!!.contains("count", ignoreCase = true))
+        }
+
+        fun checkIndexPositive(index: Int) {
+            if (index < 0) fail("Encountered negative index")
+        }
     }
 
+    @Test
+    fun indexOfOverflowSequence() {
+        assertIndexOverflow { maxIndexSequence.indexOf("j") }
+        assertIndexOverflow { maxIndexSequence.lastIndexOf("k") }
+        assertIndexOverflow { maxIndexSequence.indexOfFirst { false } }
+        assertIndexOverflow { maxIndexSequence.indexOfLast { false } }
+    }
+
+    @Test
+    fun indexOfOverflowIterable() {
+        assertIndexOverflow { maxIndexIterable.indexOf("j") }
+        assertIndexOverflow { maxIndexIterable.lastIndexOf("k") }
+        assertIndexOverflow { maxIndexIterable.indexOfFirst { false } }
+        assertIndexOverflow { maxIndexIterable.indexOfLast { false } }
+    }
+
+
+    @Test
+    fun forEachIndexedOverflow() {
+        assertIndexOverflow { maxIndexSequence.forEachIndexed { index, _ -> checkIndexPositive(index) } }
+        assertIndexOverflow { maxIndexIterable.forEachIndexed { index, _ -> checkIndexPositive(index) } }
+    }
+
+    @Test
+    fun withIndexOverflow() {
+        assertIndexOverflow { maxIndexSequence.withIndex().forEach { (index, _) -> checkIndexPositive(index) } }
+        assertIndexOverflow { maxIndexIterable.withIndex().forEach { (index, _) -> checkIndexPositive(index) } }
+    }
+
+
+    @Test
+    fun countOverflow() {
+        assertCountOverflow { repeatCounted("k").count() }
+        assertCountOverflow { repeatCounted("k").count { true } }
+        assertCountOverflow { repeatCounted("k").asIterable().count() }
+        assertCountOverflow { repeatCounted("k").asIterable().count { true } }
+    }
+
+    @Test
+    fun averageCountOverflow() {
+        assertCountOverflow { repeatCounted(1.0).average() }
+        assertCountOverflow { repeatCounted(1L).asIterable().average() }
+    }
+
+
+    private class CountingCollection<T> : AbstractMutableCollection<T>() {
+        private var _size = 0
+
+        override fun add(element: T): Boolean {
+            if (_size < 0) error("Collection is too long")
+            _size++
+            return true
+        }
+
+        override val size: Int get() = _size
+
+        override fun iterator(): MutableIterator<T> = error("not implemented")
+    }
+
+    @Test
+    fun mapIndexedOverflow() {
+        assertIndexOverflow { maxIndexSequence.mapIndexed { index, _ -> checkIndexPositive(index) }.forEach { } }
+        assertIndexOverflow { maxIndexSequence.mapIndexedTo(CountingCollection()) { index, _ -> checkIndexPositive(index) } }
+        assertIndexOverflow { maxIndexIterable.mapIndexedTo(CountingCollection()) { index, _ -> checkIndexPositive(index) } }
+    }
+
+    @Test
+    fun mapNotNullIndexedOverflow() {
+        assertIndexOverflow { maxIndexSequence.mapIndexedNotNull { index, _ -> checkIndexPositive(index) }.forEach { } }
+        assertIndexOverflow { maxIndexSequence.mapIndexedNotNullTo(CountingCollection()) { index, _ -> checkIndexPositive(index) } }
+        assertIndexOverflow { maxIndexIterable.mapIndexedNotNullTo(CountingCollection()) { index, _ -> checkIndexPositive(index) } }
+    }
+
+    @Test
+    fun filterIndexedOverflow() {
+        assertIndexOverflow { maxIndexSequence.filterIndexed { index, _ -> checkIndexPositive(index); true }.forEach { } }
+        assertIndexOverflow { maxIndexSequence.filterIndexedTo(CountingCollection()) { index, _ -> checkIndexPositive(index); true } }
+        assertIndexOverflow { maxIndexIterable.filterIndexedTo(CountingCollection()) { index, _ -> checkIndexPositive(index); true } }
+    }
+
+
+    @Test
+    fun foldIndexedOverflow() {
+        assertIndexOverflow { maxIndexSequence.foldIndexed("") { index, acc, s -> checkIndexPositive(index); s } }
+        assertIndexOverflow { maxIndexIterable.foldIndexed("") { index, acc, s -> checkIndexPositive(index); s } }
+    }
+
+    @Test
+    fun reduceIndexedOverflow() {
+        assertIndexOverflow { maxIndexSequence.reduceIndexed { index, acc, s -> checkIndexPositive(index); s } }
+        assertIndexOverflow { maxIndexIterable.reduceIndexed { index, acc, s -> checkIndexPositive(index); s } }
+    }
 
 
     @Test