~ Workaround JS BE bugs and more tests

Author: elizarov

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

@@ -51,7 +51,8 @@ public fun <T> Flow<T>.take(count: Int): Flow<T> {
     require(count > 0) { "Requested element count $count should be positive" }
     return flow {
         var consumed = 0
-        collectWhile { value ->
+        // This return is needed to work around a bug in JS BE
+        return@flow collectWhile { value ->
             emit(value)
             ++consumed < count
         }
@@ -65,7 +66,8 @@ public fun <T> Flow<T>.take(count: Int): Flow<T> {
  * See [transformWhile] for a more flexible operator.
  */
 public fun <T> Flow<T>.takeWhile(predicate: suspend (T) -> Boolean): Flow<T> = flow {
-    collectWhile { value ->
+    // This return is needed to work around a bug in JS BE
+    return@flow collectWhile { value ->
         if (predicate(value)) {
             emit(value)
             true

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

@@ -24,6 +24,27 @@ class TransformWhileTest : TestBase() {
         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(