Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Flow.transformWhile operator (#2066)
Also, most flow-truncating operators are refactored via a common internal collectWhile operator that properly uses AbortFlowException and checks for its ownership, so that we don't have to look for bugs in interactions between all those operators (and zip, too, which is also flow-truncating). But `take` operator still users a custom highly-tuned implementation. Fixes #2065 Co-authored-by: EdwarDDay <4127904+EdwarDDay@users.noreply.github.com> Co-authored-by: Louis CAD <louis.cognault@gmail.com>
- Loading branch information
3 people
committed
Jul 16, 2020
1 parent
5183b62
commit b42f986
Showing
7 changed files
with
249 additions
and
35 deletions.
There are no files selected for viewing
68 changes: 68 additions & 0 deletions
68
benchmarks/src/jmh/kotlin/benchmarks/flow/TakeWhileBenchmark.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| /* | ||
| * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license. | ||
| */ | ||
|
|
||
| @file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER") | ||
|
|
||
| package benchmarks.flow | ||
|
|
||
| import kotlinx.coroutines.* | ||
| import kotlinx.coroutines.flow.* | ||
| import kotlinx.coroutines.flow.internal.* | ||
| import org.openjdk.jmh.annotations.* | ||
| import java.util.concurrent.TimeUnit | ||
|
|
||
| @Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS) | ||
| @Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS) | ||
| @Fork(value = 1) | ||
| @BenchmarkMode(Mode.AverageTime) | ||
| @OutputTimeUnit(TimeUnit.MICROSECONDS) | ||
| @State(Scope.Benchmark) | ||
| open class TakeWhileBenchmark { | ||
| @Param("1", "10", "100", "1000") | ||
| private var size: Int = 0 | ||
|
|
||
| private suspend inline fun Flow<Long>.consume() = | ||
| filter { it % 2L != 0L } | ||
| .map { it * it }.count() | ||
|
|
||
| @Benchmark | ||
| fun baseline() = runBlocking<Int> { | ||
| (0L until size).asFlow().consume() | ||
| } | ||
|
|
||
| @Benchmark | ||
| fun takeWhileDirect() = runBlocking<Int> { | ||
| (0L..Long.MAX_VALUE).asFlow().takeWhileDirect { it < size }.consume() | ||
| } | ||
|
|
||
| @Benchmark | ||
| fun takeWhileViaCollectWhile() = runBlocking<Int> { | ||
| (0L..Long.MAX_VALUE).asFlow().takeWhileViaCollectWhile { it < size }.consume() | ||
| } | ||
|
|
||
| // Direct implementation by checking predicate and throwing AbortFlowException | ||
| private fun <T> Flow<T>.takeWhileDirect(predicate: suspend (T) -> Boolean): Flow<T> = unsafeFlow { | ||
| try { | ||
| collect { value -> | ||
| if (predicate(value)) emit(value) | ||
| else throw AbortFlowException(this) | ||
| } | ||
| } catch (e: AbortFlowException) { | ||
| e.checkOwnership(owner = this) | ||
| } | ||
| } | ||
|
|
||
| // Essentially the same code, but reusing the logic via collectWhile function | ||
| private fun <T> Flow<T>.takeWhileViaCollectWhile(predicate: suspend (T) -> Boolean): Flow<T> = unsafeFlow { | ||
| // This return is needed to work around a bug in JS BE: KT-39227 | ||
| return@unsafeFlow collectWhile { value -> | ||
| if (predicate(value)) { | ||
| emit(value) | ||
| true | ||
| } else { | ||
| false | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
70 changes: 70 additions & 0 deletions
70
kotlinx-coroutines-core/common/test/flow/operators/TransformWhileTest.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| /* | ||
| * Copyright 2016-2020 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 TransformWhileTest : TestBase() { | ||
| @Test | ||
| fun testSimple() = runTest { | ||
| val flow = (0..10).asFlow() | ||
| val expected = listOf("A", "B", "C", "D") | ||
| val actual = flow.transformWhile { value -> | ||
| when(value) { | ||
| 0 -> { emit("A"); true } | ||
| 1 -> true | ||
| 2 -> { emit("B"); emit("C"); true } | ||
| 3 -> { emit("D"); false } | ||
| else -> { expectUnreached(); false } | ||
| } | ||
| }.toList() | ||
| assertEquals(expected, actual) | ||
| } | ||
|
|
||
| @Test | ||
| fun testCancelUpstream() = runTest { | ||
| var cancelled = false | ||
| val flow = flow { | ||
| coroutineScope { | ||
| launch(start = CoroutineStart.ATOMIC) { | ||
| hang { cancelled = true } | ||
| } | ||
| emit(1) | ||
| emit(2) | ||
| emit(3) | ||
| } | ||
| } | ||
| val transformed = flow.transformWhile { | ||
| emit(it) | ||
| it < 2 | ||
| } | ||
| assertEquals(listOf(1, 2), transformed.toList()) | ||
| assertTrue(cancelled) | ||
| } | ||
|
|
||
| @Test | ||
| fun testExample() = runTest { | ||
| val source = listOf( | ||
| DownloadProgress(0), | ||
| DownloadProgress(50), | ||
| DownloadProgress(100), | ||
| DownloadProgress(147) | ||
| ) | ||
| val expected = source.subList(0, 3) | ||
| val actual = source.asFlow().completeWhenDone().toList() | ||
| assertEquals(expected, actual) | ||
| } | ||
|
|
||
| private fun Flow<DownloadProgress>.completeWhenDone(): Flow<DownloadProgress> = | ||
| transformWhile { progress -> | ||
| emit(progress) // always emit progress | ||
| !progress.isDone() // continue while download is not done | ||
| } | ||
|
|
||
| private data class DownloadProgress(val percent: Int) { | ||
| fun isDone() = percent >= 100 | ||
| } | ||
| } |