Introduce ReceiveChannel.consumeAsFlow and FlowCollector.emitAll(chan)

* This is a consuming conversion -- the resulting flow can be collected
  just once and the channel is closed after the first collect.
* The implementation is made efficient via emitAll extension.
* Experimental FlowCollector.emitAll extension is introduced.
* It is based on the (internal) Channel.receiveOrClose
  and ensures that the reference to the last emitted value is
  not retained (does not leak).

Fixes #1340
Fixes #1333

Author: elizarov

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

@@ -848,6 +848,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun combineLatest (Lkotlinx/coroutines/flow/Flow;[Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function0;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final synthetic fun combineLatest (Lkotlinx/coroutines/flow/Flow;[Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun conflate (Lkotlinx/coroutines/flow/Flow;)Lkotlinx/coroutines/flow/Flow;
+	public static final fun consumeAsFlow (Lkotlinx/coroutines/channels/ReceiveChannel;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun count (Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun count (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun debounce (Lkotlinx/coroutines/flow/Flow;J)Lkotlinx/coroutines/flow/Flow;
@@ -857,6 +858,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static final fun distinctUntilChangedBy (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function1;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun drop (Lkotlinx/coroutines/flow/Flow;I)Lkotlinx/coroutines/flow/Flow;
 	public static final fun dropWhile (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
+	public static final fun emitAll (Lkotlinx/coroutines/flow/FlowCollector;Lkotlinx/coroutines/channels/ReceiveChannel;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun emitAll (Lkotlinx/coroutines/flow/FlowCollector;Lkotlinx/coroutines/flow/Flow;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun emptyFlow ()Lkotlinx/coroutines/flow/Flow;
 	public static final fun filter (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;

kotlinx-coroutines-core/common/src/channels/AbstractChannel.kt

@@ -6,6 +6,7 @@ package kotlinx.coroutines.channels
 
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
 import kotlinx.coroutines.selects.*

kotlinx-coroutines-core/common/src/channels/Channel.kt

@@ -7,12 +7,13 @@
 package kotlinx.coroutines.channels
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.Channel.Factory.BUFFERED
+import kotlinx.coroutines.channels.Channel.Factory.CHANNEL_DEFAULT_CAPACITY
 import kotlinx.coroutines.channels.Channel.Factory.CONFLATED
 import kotlinx.coroutines.channels.Channel.Factory.RENDEZVOUS
 import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
-import kotlinx.coroutines.channels.Channel.Factory.BUFFERED
-import kotlinx.coroutines.channels.Channel.Factory.CHANNEL_DEFAULT_CAPACITY
-import kotlinx.coroutines.internal.systemProp
+import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.selects.*
 import kotlin.jvm.*
 import kotlin.internal.*

kotlinx-coroutines-core/common/src/flow/Channels.kt

@@ -7,14 +7,81 @@
 
 package kotlinx.coroutines.flow
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
-import kotlinx.coroutines.channels.Channel.Factory.CONFLATED
-import kotlinx.coroutines.channels.Channel.Factory.BUFFERED
-import kotlinx.coroutines.channels.Channel.Factory.OPTIONAL_CHANNEL
 import kotlinx.coroutines.flow.internal.*
-import kotlin.coroutines.*
 import kotlin.jvm.*
+import kotlinx.coroutines.flow.unsafeFlow as flow
+
+/**
+ * Emits all elements from the given [channel] to this flow collector and [cancels][cancel] (consumes)
+ * the channel afterwards. If you need to iterate over the channel without consuming it,
+ * a regular `for` loop should be used instead.
+ *
+ * This function provides a more efficient shorthand for `channel.consumeEach { value -> emit(value) }`.
+ * See [consumeEach][ReceiveChannel.consumeEach].
+ */
+@ExperimentalCoroutinesApi
+public suspend fun <T> FlowCollector<T>.emitAll(channel: ReceiveChannel<T>) {
+    // Manually inlined "consumeEach" implementation that does not use iterator but works via "receiveOrClosed".
+    // It has smaller and more efficient spilled state which also allows to implement a manual kludge to
+    // fix retention of the last emitted value.
+    // See https://youtrack.jetbrains.com/issue/KT-16222
+    // See https://github.com/Kotlin/kotlinx.coroutines/issues/1333
+    var cause: Throwable? = null
+    try {
+        while (true) {
+            // :KLUDGE: This "run" call is resolved to an extension function "run" and forces the size of
+            // spilled state to increase by an additional slot, so there are 4 object local variables spilled here
+            // which makes the size of spill state equal to the 4 slots that are spilled around subsequent "emit"
+            // call, ensuring that the previously emitted value is not retained in the state while receiving
+            // the next one.
+            //     L$0 <- this
+            //     L$1 <- channel
+            //     L$2 <- cause
+            //     L$3 <- this$run (actually equal to this)
+            val result = run { channel.receiveOrClosed() }
+            if (result.isClosed) {
+                result.closeCause?.let { throw it }
+                break // returns normally when result.closeCause == null
+            }
+            // result is spilled here to the coroutine state and retained after the call, even though
+            // it is not actually needed in the next loop iteration.
+            //     L$0 <- this
+            //     L$1 <- channel
+            //     L$2 <- cause
+            //     L$3 <- result
+            emit(result.value)
+        }
+    } catch (e: Throwable) {
+        cause = e
+        throw e
+    } finally {
+        channel.cancelConsumed(cause)
+    }
+}
+
+/**
+ * Represents the given receive channel as a hot flow and [consumes][ReceiveChannel.consume] the channel
+ * on the first collection from this flow. The resulting flow can be collected just once and throws
+ * [IllegalStateException] when trying to collect it more than once.
+ *
+ * ### Cancellation semantics
+ * 1) Flow consumer is cancelled when the original channel is cancelled.
+ * 2) Flow consumer completes normally when the original channel completes (~is closed) normally.
+ * 3) If the flow consumer fails with an exception, channel is cancelled.
+ *
+ */
+@FlowPreview
+public fun <T> ReceiveChannel<T>.consumeAsFlow(): Flow<T> = object : Flow<T> {
+    val collected = atomic(false)
+
+    override suspend fun collect(collector: FlowCollector<T>) {
+        check(!collected.getAndSet(true)) { "ReceiveChannel.consumeAsFlow can be collected just once" }
+        collector.emitAll(this@consumeAsFlow)
+    }
+}
 
 /**
  * Represents the given broadcast channel as a hot flow.
@@ -27,10 +94,7 @@ import kotlin.jvm.*
  */
 @FlowPreview
 public fun <T> BroadcastChannel<T>.asFlow(): Flow<T> = flow {
-    val subscription = openSubscription()
-    subscription.consumeEach { value ->
-        emit(value)
-    }
+    emitAll(openSubscription())
 }
 
 /**

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

@@ -72,8 +72,7 @@ public abstract class ChannelFlow<T>(
 
     override suspend fun collect(collector: FlowCollector<T>) =
         coroutineScope {
-            val channel = produceImpl(this)
-            channel.consumeEach { collector.emit(it) }
+            collector.emitAll(produceImpl(this))
         }
 
     // debug toString

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

@@ -61,7 +61,7 @@ private class ChannelViaBroadcast<E>(
     override suspend fun receiveOrClosed(): ValueOrClosed<E> = sub.receiveOrClosed()
     override fun poll(): E? = sub.poll()
     override fun iterator(): ChannelIterator<E> = sub.iterator()
-    
+
     override fun cancel(cause: CancellationException?) = sub.cancel(cause)
 
     // implementing hidden method anyway, so can cast to an internal class