Flow firstOrNull support (#1869)

* Flow firstOrNull support

Co-authored-by: Bradyn Poulsen <[email protected]>

Author: qwwdfsad

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -907,6 +907,8 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun filterNotNull (Lkotlinx/coroutines/flow/Flow;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun first (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun first (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun firstOrNull (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun firstOrNull (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun flatMap (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun flatMapConcat (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun flatMapLatest (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;

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

@@ -77,7 +77,6 @@ public suspend fun <T: Any> Flow<T>.singleOrNull(): T? {
         if (result != null) error("Expected only one element")
         result = value
     }
-
     return result
 }
 
@@ -120,3 +119,39 @@ public suspend fun <T> Flow<T>.first(predicate: suspend (T) -> Boolean): T {
     if (result === NULL) throw NoSuchElementException("Expected at least one element matching the predicate $predicate")
     return result as T
 }
+
+/**
+ * The terminal operator that returns the first element emitted by the flow and then cancels flow's collection.
+ * Returns `null` if the flow was empty.
+ */
+public suspend fun <T : Any> Flow<T>.firstOrNull(): T? {
+    var result: T? = null
+    try {
+        collect { value ->
+            result = value
+            throw AbortFlowException(NopCollector)
+        }
+    } catch (e: AbortFlowException) {
+        // Do nothing
+    }
+    return result
+}
+
+/**
+ *  The terminal operator that returns the first element emitted by the flow matching the given [predicate] and then cancels flow's collection.
+ * Returns `null` if the flow did not contain an element matching the [predicate].
+ */
+public suspend fun <T : Any> Flow<T>.firstOrNull(predicate: suspend (T) -> Boolean): T? {
+    var result: T? = null
+    try {
+        collect { value ->
+            if (predicate(value)) {
+                result = value
+                throw AbortFlowException(NopCollector)
+            }
+        }
+    } catch (e: AbortFlowException) {
+        // Do nothing
+    }
+    return result
+}

kotlinx-coroutines-core/common/test/AsyncTest.kt

@@ -210,12 +210,6 @@ class AsyncTest : TestBase() {
         finish(13)
     }
 
-    class BadClass {
-        override fun equals(other: Any?): Boolean = error("equals")
-        override fun hashCode(): Int = error("hashCode")
-        override fun toString(): String = error("toString")
-    }
-
     @Test
     fun testDeferBadClass() = runTest {
         val bad = BadClass()

kotlinx-coroutines-core/common/test/TestBase.common.kt

@@ -80,3 +80,8 @@ public fun wrapperDispatcher(context: CoroutineContext): CoroutineContext {
 
 public suspend fun wrapperDispatcher(): CoroutineContext = wrapperDispatcher(coroutineContext)
 
+class BadClass {
+    override fun equals(other: Any?): Boolean = error("equals")
+    override fun hashCode(): Int = error("hashCode")
+    override fun toString(): String = error("toString")
+}

kotlinx-coroutines-core/common/test/WithTimeoutOrNullTest.kt

@@ -152,12 +152,6 @@ class WithTimeoutOrNullTest : TestBase() {
         assertSame(bad, result)
     }
 
-    class BadClass {
-        override fun equals(other: Any?): Boolean = error("Should not be called")
-        override fun hashCode(): Int = error("Should not be called")
-        override fun toString(): String = error("Should not be called")
-    }
-
     @Test
     fun testNullOnTimeout() = runTest {
         expect(1)

kotlinx-coroutines-core/common/test/WithTimeoutTest.kt

@@ -107,12 +107,6 @@ class WithTimeoutTest : TestBase() {
         assertSame(bad, result)
     }
 
-    class BadClass {
-        override fun equals(other: Any?): Boolean = error("Should not be called")
-        override fun hashCode(): Int = error("Should not be called")
-        override fun toString(): String = error("Should not be called")
-    }
-
     @Test
     fun testExceptionOnTimeout() = runTest {
         expect(1)

kotlinx-coroutines-core/common/test/channels/RendezvousChannelTest.kt

@@ -241,12 +241,6 @@ class RendezvousChannelTest : TestBase() {
         finish(12)
     }
 
-    class BadClass {
-        override fun equals(other: Any?): Boolean = error("equals")
-        override fun hashCode(): Int = error("hashCode")
-        override fun toString(): String = error("toString")
-    }
-
     @Test
     fun testProduceBadClass() = runTest {
         val bad = BadClass()

kotlinx-coroutines-core/common/test/flow/terminal/FirstTest.kt

@@ -83,4 +83,81 @@ class FirstTest : TestBase() {
         assertEquals(1, flow.first())
         finish(2)
     }
+
+    @Test
+    fun testFirstOrNull() = runTest {
+        val flow = flowOf(1, 2, 3)
+        assertEquals(1, flow.firstOrNull())
+    }
+
+    @Test
+    fun testFirstOrNullWithPredicate() = runTest {
+        val flow = flowOf(1, 2, 3)
+        assertEquals(1, flow.firstOrNull { it > 0 })
+        assertEquals(2, flow.firstOrNull { it > 1 })
+        assertNull(flow.firstOrNull { it > 3 })
+    }
+
+    @Test
+    fun testFirstOrNullCancellation() = runTest {
+        val latch = Channel<Unit>()
+        val flow = flow {
+            coroutineScope {
+                launch {
+                    latch.send(Unit)
+                    hang { expect(1) }
+                }
+                emit(1)
+                emit(2)
+            }
+        }
+
+
+        val result = flow.firstOrNull {
+            latch.receive()
+            true
+        }
+        assertEquals(1, result)
+        finish(2)
+    }
+
+    @Test
+    fun testFirstOrNullWithEmptyFlow() = runTest {
+        assertNull(emptyFlow<Int>().firstOrNull())
+        assertNull(emptyFlow<Int>().firstOrNull { true })
+    }
+
+    @Test
+    fun testFirstOrNullWhenErrorCancelsUpstream() = runTest {
+        val latch = Channel<Unit>()
+        val flow = flow {
+            coroutineScope {
+                launch {
+                    latch.send(Unit)
+                    hang { expect(1) }
+                }
+                emit(1)
+            }
+        }
+
+        assertFailsWith<TestException> {
+            flow.firstOrNull {
+                latch.receive()
+                throw TestException()
+            }
+        }
+
+        assertEquals(1, flow.firstOrNull())
+        finish(2)
+    }
+
+    @Test
+    fun testBadClass() = runTest {
+        val instance = BadClass()
+        val flow = flowOf(instance)
+        assertSame(instance, flow.first())
+        assertSame(instance, flow.firstOrNull())
+        assertSame(instance, flow.first { true })
+        assertSame(instance, flow.firstOrNull { true })
+    }
 }

kotlinx-coroutines-core/common/test/flow/terminal/SingleTest.kt

@@ -17,7 +17,6 @@ class SingleTest : TestBase() {
 
         assertEquals(239L, flow.single())
         assertEquals(239L, flow.singleOrNull())
-
     }
 
     @Test
@@ -63,4 +62,12 @@ class SingleTest : TestBase() {
         assertNull(flowOf<Int?>(null).single())
         assertFailsWith<NoSuchElementException> { flowOf<Int?>().single() }
     }
+
+    @Test
+    fun testBadClass() = runTest {
+        val instance = BadClass()
+        val flow = flowOf(instance)
+        assertSame(instance, flow.single())
+        assertSame(instance, flow.singleOrNull())
+    }
 }