withIndex and collectIndexed operators

Fixes #1247

Author: qwwdfsad

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt

@@ -824,6 +824,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun channelFlow (Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun collect (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun collect (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun collectIndexed (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function3;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun combineLatest (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function3;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun combineLatest (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function4;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun combineLatest (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function5;)Lkotlinx/coroutines/flow/Flow;
@@ -921,6 +922,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun unsafeFlow (Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun unsafeTransform (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function3;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun withContext (Lkotlinx/coroutines/flow/FlowCollector;Lkotlin/coroutines/CoroutineContext;Lkotlin/jvm/functions/Function1;)V
+	public static final fun withIndex (Lkotlinx/coroutines/flow/Flow;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun zip (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function3;)Lkotlinx/coroutines/flow/Flow;
 }
 
@@ -939,6 +941,10 @@ public abstract class kotlinx/coroutines/flow/internal/ChannelFlow : kotlinx/cor
 	public static synthetic fun update$default (Lkotlinx/coroutines/flow/internal/ChannelFlow;Lkotlin/coroutines/CoroutineContext;IILjava/lang/Object;)Lkotlinx/coroutines/flow/internal/ChannelFlow;
 }
 
+public final class kotlinx/coroutines/flow/internal/FlowExceptions_commonKt {
+	public static final fun checkIndexOverflow (I)I
+}
+
 public final class kotlinx/coroutines/flow/internal/SafeCollector : kotlinx/coroutines/flow/FlowCollector {
 	public fun <init> (Lkotlinx/coroutines/flow/FlowCollector;Lkotlin/coroutines/CoroutineContext;)V
 	public fun emit (Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;

kotlinx-coroutines-core/common/src/flow/internal/FlowExceptions.common.kt

@@ -16,3 +16,12 @@ internal expect class AbortFlowException() : CancellationException
  * Exception used to cancel child of [scopedFlow] without cancelling the whole scope.
  */
 internal expect class ChildCancelledException() : CancellationException
+
+@Suppress("NOTHING_TO_INLINE")
+@PublishedApi
+internal inline fun checkIndexOverflow(index: Int): Int {
+    if (index < 0) {
+        throw ArithmeticException("Index overflow has happened")
+    }
+    return index
+}

kotlinx-coroutines-core/common/src/flow/operators/Limit.kt

@@ -2,7 +2,6 @@
  * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
-
 @file:JvmMultifileClass
 @file:JvmName("FlowKt")
 

kotlinx-coroutines-core/common/src/flow/operators/Transform.kt

@@ -62,6 +62,17 @@ public inline fun <T, R: Any> Flow<T>.mapNotNull(crossinline transform: suspend
     return@transform emit(transformed)
 }
 
+/**
+ * Returns a flow that wraps each element into [IndexedValue], containing value and its index (starting from zero).
+ */
+@ExperimentalCoroutinesApi
+public fun <T> Flow<T>.withIndex(): Flow<IndexedValue<T>> = flow {
+    var index = 0
+    collect { value ->
+        emit(IndexedValue(checkIndexOverflow(index++), value))
+    }
+}
+
 /**
  * Returns a flow which performs the given [action] on each value of the original flow.
  */

kotlinx-coroutines-core/common/src/flow/terminal/Collect.kt

@@ -75,6 +75,19 @@ public suspend inline fun <T> Flow<T>.collect(crossinline action: suspend (value
         override suspend fun emit(value: T) = action(value)
     })
 
+/**
+ * Terminal flow operator that collects the given flow with a provided [action] that takes the index of an element (zero-based) and the element.
+ * If any exception occurs during collect or in the provided flow, this exception is rethrown from this method.
+ *
+ * See also [collect] and [withIndex].
+ */
+@ExperimentalCoroutinesApi
+public suspend inline fun <T> Flow<T>.collectIndexed(crossinline action: suspend (index: Int, value: T) -> Unit): Unit =
+    collect(object : FlowCollector<T> {
+        private var index = 0
+        override suspend fun emit(value: T) = action(checkIndexOverflow(index++), value)
+    })
+
 /**
  * Collects all the values from the given [flow] and emits them to the collector.
  * 

kotlinx-coroutines-core/common/test/flow/operators/IndexedTest.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.flow
+
+import kotlinx.coroutines.*
+import kotlin.test.*
+
+class IndexedTest : TestBase() {
+
+    @Test
+    fun testWithIndex() = runTest {
+        val flow = flowOf(3, 2, 1).withIndex()
+        assertEquals(listOf(IndexedValue(0, 3), IndexedValue(1, 2), IndexedValue(2, 1)), flow.toList())
+    }
+
+    @Test
+    fun testWithIndexEmpty() = runTest {
+        val flow = emptyFlow<Int>().withIndex()
+        assertEquals(emptyList(), flow.toList())
+    }
+
+    @Test
+    fun testCollectIndexed() = runTest {
+        val result = ArrayList<IndexedValue<Long>>()
+        flowOf(3L, 2L, 1L).collectIndexed { index, value ->
+            result.add(IndexedValue(index, value))
+        }
+        assertEquals(listOf(IndexedValue(0, 3L), IndexedValue(1, 2L), IndexedValue(2, 1L)), result)
+    }
+
+    @Test
+    fun testCollectIndexedEmptyFlow() = runTest {
+        val flow = flow<Int> {
+            expect(1)
+        }
+
+        flow.collectIndexed { _, _ ->
+            expectUnreached()
+        }
+
+        finish(2)
+    }
+}