Add onEachIndexed similar to forEachIndexed #KT-37161

libraries/stdlib/test/collections/ArraysTest.kt

@@ -1441,6 +1441,23 @@ class ArraysTest {
         assertEquals(listOf('1', '2', '3'), mutableListOf<Char>().apply { charArrayOf('1', '2', '3').onEach { add(it) } })
     }
 
+    @Test fun onEachIndexed() {
+        assertEquals(listOf(1, 3, 5), mutableListOf<Int>().apply { intArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i + e) } })
+        assertEquals(listOf(1, 3, 5), mutableListOf<Int>().apply { byteArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i + e) } })
+        assertEquals(listOf(1, 3, 5), mutableListOf<Int>().apply { shortArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i + e) } })
+        assertEquals(listOf<Long>(1, 3, 5), mutableListOf<Long>().apply { longArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i + e) } })
+        assertEquals(listOf(1f, 3f, 5f), mutableListOf<Float>().apply { floatArrayOf(1f, 2f, 3f).onEachIndexed { i, e -> add(i + e) } })
+        assertEquals(listOf(1.0, 3.0, 5.0), mutableListOf<Double>().apply { doubleArrayOf(1.0, 2.0, 3.0).onEachIndexed { i, e -> add(i + e) } })
+        assertEquals(listOf(true, false, true), mutableListOf<Boolean>().apply { booleanArrayOf(true, false, false).onEachIndexed { i, e -> add(i % 2 == 0 || e) } })
+        assertEquals(listOf('1', '3', '5'), mutableListOf<Char>().apply { charArrayOf('1', '2', '3').onEachIndexed { i, e -> add(e + i) } })
+        assertEquals(listOf("a0", "b1", "c2"), mutableListOf<String>().apply { arrayOf("a", "b", "c").onEachIndexed { i, e -> add(e + i) } })
+
+        val empty = arrayOf<Int>()
+        assertSame(empty, empty.onEachIndexed { i, e -> fail("Should be unreachable: $i, $e") })
+        val nonEmpty = longArrayOf(1, 2, 3)
+        assertSame(nonEmpty, nonEmpty.onEachIndexed { _, _ -> })
+    }
+
     @Test fun drop() {
         expect(listOf(1), { intArrayOf(1).drop(0) })
         expect(listOf(), { intArrayOf().drop(1) })

libraries/stdlib/test/collections/IterableTests.kt

@@ -313,6 +313,17 @@ abstract class IterableTests<T : Iterable<String>>(val createFrom: (Array<out St
         assertStaticTypeIs<ArrayList<Int>>(arrayListOf(1, 2, 3).onEach { })
     }
 
+    @Test
+    fun onEachIndexed() {
+        var count = 0
+        val newData = data.onEachIndexed { i, e -> count += i + e.length }
+        assertEquals(7, count)
+        assertSame(data, newData)
+
+        // static types test
+        assertStaticTypeIs<ArrayList<Int>>(arrayListOf(1, 2, 3).onEachIndexed { _, _ -> })
+    }
+
     @Test
     fun contains() {
         assertTrue(data.contains("foo"))

libraries/stdlib/test/collections/MapTest.kt

@@ -128,6 +128,20 @@ class MapTest {
         )
     }
 
+    @Test
+    fun onEachIndexed() {
+        val map = mutableMapOf("beverage" to "beer", "location" to "Mells")
+        val result = StringBuilder()
+        val newMap = map.onEachIndexed { i, e -> result.append(i + 1).append('.').append(e.key).append("=").append(e.value).append(";") }
+        assertEquals("1.beverage=beer;2.location=Mells;", result.toString())
+        assertTrue(map === newMap)
+
+        // static types test
+        assertStaticTypeIs<HashMap<String, String>>(
+            hashMapOf("a" to "b").onEachIndexed { _, _ -> }
+        )
+    }
+
     @Test fun stream() {
         val map = mapOf("beverage" to "beer", "location" to "Mells", "name" to "James")
         val named = map.asSequence().filter { it.key == "name" }.single()

libraries/stdlib/test/collections/SequenceTest.kt

@@ -126,6 +126,21 @@ public class SequenceTest {
         assertEquals(sum, count)
     }
 
+    @Test
+    fun onEachIndexed() {
+        var count = 0
+        val data = sequenceOf("foo", "bar")
+        val newData = data.onEachIndexed { i, e -> count += i + e.length }
+        assertNotSame(data, newData)
+        assertEquals(0, count, "onEachIndex should be executed lazily")
+
+        data.forEach {  }
+        assertEquals(0, count, "onEachIndex should be executed only when resulting sequence is iterated")
+
+        val sum = newData.foldIndexed(0) { i, acc, e -> acc + i + e.length }
+        assertEquals(sum, count)
+    }
+
 
     @Test fun filterAndTakeWhileExtractTheElementsWithinRange() {
         assertEquals(listOf(144, 233, 377, 610, 987), fibonacci().filter { it > 100 }.takeWhile { it < 1000 }.toList())

libraries/stdlib/test/collections/UnsignedArraysTest.kt

@@ -994,6 +994,21 @@ class UnsignedArraysTest {
         assertEquals(listOf<ULong>(1, 2, 3), mutableListOf<ULong>().apply { ulongArrayOf(1, 2, 3).onEach { add(it) } })
     }
 
+    @Test
+    fun onEachIndexed() {
+        assertEquals(listOf<UInt>(1, 3, 5), mutableListOf<UInt>().apply { uintArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i.toUInt() + e) } })
+        assertEquals(listOf<UInt>(1, 3, 5), mutableListOf<UInt>().apply { ubyteArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i.toUByte() + e) } })
+        assertEquals(listOf<UInt>(1, 3, 5), mutableListOf<UInt>().apply { ushortArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i.toUShort() + e) } })
+        assertEquals(listOf<ULong>(1, 3, 5), mutableListOf<ULong>().apply { ulongArrayOf(1, 2, 3).onEachIndexed { i, e -> add(i.toULong() + e) } })
+
+        val empty = arrayOf<UInt>()
+        assertSame(empty, empty.onEachIndexed { i, e -> fail("Should be unreachable: $i, $e") })
+
+        // Identity equality for arguments of types ULongArray and ULongArray is forbidden
+//        val nonEmpty = ulongArrayOf(1, 2, 3)
+//        assertSame(nonEmpty, nonEmpty.onEachIndexed { _, _ -> })
+    }
+
     @Test
     fun drop() {
         expect(listOf(1.toUByte())) { ubyteArrayOf(1).drop(0) }

libraries/stdlib/test/text/StringTest.kt

@@ -948,6 +948,18 @@ class StringTest {
         assertStaticTypeIs<StringBuilder>(result.onEach { })
     }
 
+    @Test
+    fun onEachIndexed() = withOneCharSequenceArg("abcd") { data ->
+        val result = StringBuilder()
+        val newData = data.onEachIndexed { i, e -> result.append(e + i) }
+        assertEquals("aceg", result.toString())
+        assertSame(data, newData)
+
+        // static types test
+        assertStaticTypeIs<String>("x".onEachIndexed { _, _ -> })
+        assertStaticTypeIs<StringBuilder>(result.onEachIndexed { _, _ -> })
+    }
+
 
     @Test fun filter() {
         assertEquals("acdca", ("abcdcba").filter { !it.equals('b') })