Provide flatMapIndexed operation

- similar to flatMap, but transform function takes index and element

#KT-36894

libraries/stdlib/test/collections/ArraysTest.kt

@@ -1786,6 +1786,15 @@ class ArraysTest {
         assertEquals(expected, result2)
     }
 
+    @Test fun flatMapIndexed() {
+        val data = arrayOf(arrayOf(1, 2, 3), arrayOf(4, 5, 6))
+        val result1 = data.flatMapIndexed { index, arr -> arr.asList().subList(0, index + 1) }
+        val result2 = data.flatMapIndexed { index, arr -> arr.asSequence().take(index + 1) }
+        val expected = listOf(1, 4, 5)
+        assertEquals(expected, result1)
+        assertEquals(expected, result2)
+    }
+
     @Test fun flattenArray() {
         val arr1: Array<Array<Int>> = arrayOf(arrayOf(1, 2, 3), arrayOf(4, 5, 6))
         val arr2: Array<out Array<Int>> = arr1

libraries/stdlib/test/collections/CollectionTest.kt

@@ -60,6 +60,16 @@ class CollectionTest {
         assertEquals(expected, result2)
     }
 
+    @Test fun flatMapIndexed() {
+        val source = listOf(null, "foo", "bar")
+        val result1 = source.flatMapIndexed { index, it -> it.orEmpty().take(index + 1).asSequence() }
+        val result2 = source.flatMapIndexed { index, it -> it.orEmpty().take(index + 1).asIterable() }
+
+        val expected = "fobar".toList()
+        assertEquals(expected, result1)
+        assertEquals(expected, result2)
+    }
+
     /*
     @Test fun mapNotNull() {
         val data = listOf(null, "foo", null, "bar")

libraries/stdlib/test/collections/SequenceTest.kt

@@ -592,6 +592,14 @@ public class SequenceTest {
         assertEquals(expected, result2.toList())
     }
 
+    @Test fun flatMapIndexed() {
+        val result1 = sequenceOf(1, 2).flatMapIndexed { index, v -> (0..v + index).asSequence() }
+        val result2 = sequenceOf(1, 2).flatMapIndexed { index, v -> 0..v + index }
+        val expected = listOf(0, 1, 0, 1, 2, 3)
+        assertEquals(expected, result1.toList())
+        assertEquals(expected, result2.toList())
+    }
+
     @Test fun flatten() {
         val expected = listOf(0, 1, 0, 1, 2)
 

libraries/stdlib/test/collections/UnsignedArraysTest.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -981,6 +981,14 @@ class UnsignedArraysTest {
         assertEquals(listOf(), ulongArrayOf(1, 2, 3).flatMap { listOf<ULong>() })
     }
 
+    @Test
+    fun flatMapIndexed() {
+        assertEquals(listOf(), ubyteArrayOf().flatMapIndexed { index, _ -> listOf(index) })
+        assertEquals(listOf<UShort>(2, 3, 3), ushortArrayOf(1, 2, 3).flatMapIndexed { index, e -> List(index) { e } })
+        assertEquals(listOf<UInt>(2, 2, 3, 3, 3, 3), uintArrayOf(1, 2, 3).flatMapIndexed { index, e -> List(index * 2) { e } })
+        assertEquals(listOf(), ulongArrayOf(1, 2, 3).flatMapIndexed { _, _ -> listOf<ULong>() })
+    }
+
     @Test
     fun withIndex() {
         fun <T> assertIterableContentEquals(expected: Iterable<T>, actual: Iterable<T>) {

libraries/stdlib/test/text/StringTest.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -1242,6 +1242,11 @@ class StringTest {
         assertEquals("a ab abc abcd ".toList(), result)
     }
 
+    @Test fun flatMapIndexed() = withOneCharSequenceArg("bdfi") { data ->
+        val result = data.flatMapIndexed { index, c -> ('a'..c).drop(index) + ' ' }
+        assertEquals("ab bcd cdef defghi ".toList(), result)
+    }
+
     @Test fun fold() = withOneCharSequenceArg { arg1 ->
         // calculate number of digits in the string
         val data = arg1("a1b2c3def")