withIndex and collectIndexed operators

Fixes #1247

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

@@ -845,6 +845,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;
@@ -944,6 +945,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;
 }
 
@@ -962,6 +964,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)
+ }
+}