Support BufferedChannel

Before, `buffer` operations and such were only supported partially, and in some cases async stack traces were working only in 50% collects. For example, when several flows are merged with `flattenMerge` and emit values simultaneously.

This change seems large, changing a lot of lines in BufferedChannel.kt, but most of them are effectively refactoring (propagation of a wrapped value).

Author: AlexVanGogen

IntelliJ-patches.md

@@ -64,7 +64,7 @@ The key for MutableStateFlow is the element itself. For MutableSharedFlow, the e
 Most of the operators applicable to flows (such as `map`, `scan`, `debounce`, `buffer`) are supported. As some of them use an intermediary flow inside, the transferred values are wrapped and unwrapped the same way as in MutableSharedFlow.
 It means there may be all-library async stack traces between a stack trace containing `emit` and a stack trace containing `collect`.
 
-There is no support yet for many operators that heavily use `Channel`s inside (such as `timeout`), as well as for functions that convert flows to channels and vice versa (such as `produceIn`).
+There is no support yet for many operators that use `Select` inside (such as `timeout`).
 
 ### API
 
@@ -75,6 +75,7 @@ Some logic related to instrumentation was extracted to separate methods so that
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.wrapInternalDebuggerCapture` -- wraps passed arguments into a `FlowValueWrapperInternal`; only used after transformation.
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.unwrapInternal` -- returns passed argument by default; the agent instruments it to call `unwrapInternalDebuggerCapture` instead
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.unwrapInternalDebuggerCapture` -- unwraps passed argument so it returns the original value; only used after transformation
+- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.unwrapTyped` -- utility function served to ease casting to a real underlying type
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.emitInternal(FlowCollector, value)` -- common insertion point for a debugger agent; simplifies instrumentation; the value is always being unwrapped inside
 
 One internal method was added to `BufferedChannelIterator`: `nextInternal` -- same as `next` but may return a wrapped value. It should only be used with a function that is capable of unwrapping the value (see `BufferedChannel.emitAll` and `BufferedChannelIterator.next`), so there's a guarantee a wrapped value will always unwrap before emitting.
@@ -83,3 +84,5 @@ Why not just let `next` return a maybe wrapped value? That's because it is heavi
 
 One public method was added to support `buffer` and operators that use it inside:
 - `ReceiveChannel.emitAll`. It encapsulates emitting values in `FlowCollector.emitAllImpl` and has a special implementation in `BufferedChannel`.
+
+Changes were made to lambda parameter `onElementRetrieved` in `BufferedChannel<E>` methods: now they accept `Any?` instead of `E` because now they may be given a wrapped value.

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

@@ -665,7 +665,7 @@ internal open class BufferedChannel<E>(
     protected open fun onReceiveDequeued() {}
 
     override suspend fun receive(): E =
-        receiveImpl( // <-- this is an inline function
+        receiveImpl<E>( // <-- this is an inline function
             // Do not create a continuation until it is required;
             // it is created later via [onNoWaiterSuspend], if needed.
             waiter = null,
@@ -674,7 +674,7 @@ internal open class BufferedChannel<E>(
             // Also, inform `BufferedChannel` extensions that
             // synchronization of this receive operation is completed.
             onElementRetrieved = { element ->
-                return element
+                return unwrapTyped<E>(element)
             },
             // As no waiter is provided, suspension is impossible.
             onSuspend = { _, _, _ -> error("unexpected") },
@@ -683,7 +683,7 @@ internal open class BufferedChannel<E>(
             // If `receive()` decides to suspend, the corresponding
             // `suspend` function that creates a continuation is called.
             // The tail-call optimization is applied here.
-            onNoWaiterSuspend = { segm, i, r -> receiveOnNoWaiterSuspend(segm, i, r) }
+            onNoWaiterSuspend = { segm, i, r -> unwrapTyped(receiveOnNoWaiterSuspend(segm, i, r)) }
         )
 
     private suspend fun receiveOnNoWaiterSuspend(
@@ -706,7 +706,7 @@ internal open class BufferedChannel<E>(
             // not dispatched yet. In case `onUndeliveredElement` is
             // specified, we need to invoke it in the latter case.
             onElementRetrieved = { element ->
-                val onCancellation = onUndeliveredElement?.bindCancellationFun(element, cont.context)
+                val onCancellation = onUndeliveredElement?.bindCancellationFun(unwrapTyped(element), cont.context)
                 cont.resume(element, onCancellation)
             },
             onClosed = { onClosedReceiveOnNoWaiterSuspend(cont) },
@@ -733,11 +733,11 @@ internal open class BufferedChannel<E>(
         receiveImpl( // <-- this is an inline function
             waiter = null,
             onElementRetrieved = { element ->
-                success(element)
+                success(unwrapTyped(element))
             },
             onSuspend = { _, _, _ -> error("unexpected") },
             onClosed = { closed(closeCause) },
-            onNoWaiterSuspend = { segm, i, r -> receiveCatchingOnNoWaiterSuspend(segm, i, r) }
+            onNoWaiterSuspend = { segm, i, r -> unwrapTyped(receiveCatchingOnNoWaiterSuspend(segm, i, r)) }
         )
 
     private suspend fun receiveCatchingOnNoWaiterSuspend(
@@ -750,7 +750,8 @@ internal open class BufferedChannel<E>(
             segment, index, r,
             waiter = waiter,
             onElementRetrieved = { element ->
-                cont.resume(success(element), onUndeliveredElement?.bindCancellationFun(element, cont.context))
+                val unwrapped = unwrapTyped<E>(element)
+                cont.resume(success(unwrapped), onUndeliveredElement?.bindCancellationFun(unwrapped, cont.context))
             },
             onClosed = { onClosedReceiveCatchingOnNoWaiterSuspend(cont) }
         )
@@ -783,7 +784,7 @@ internal open class BufferedChannel<E>(
             // Store an already interrupted receiver in case of suspension.
             waiter = INTERRUPTED_RCV,
             // Finish when an element is successfully retrieved.
-            onElementRetrieved = { element -> success(element) },
+            onElementRetrieved = { element -> success(unwrapTyped(element)) },
             // On suspension, the `INTERRUPTED_RCV` token has been
             // installed, and this `tryReceive()` must fail.
             onSuspend = { segm, _, globalIndex ->
@@ -847,7 +848,7 @@ internal open class BufferedChannel<E>(
                     // Clean the reference to the previous segment.
                     segment.cleanPrev()
                     @Suppress("UNCHECKED_CAST")
-                    onUndeliveredElement?.callUndeliveredElementCatchingException(updCellResult as E)?.let { throw it }
+                    onUndeliveredElement?.callUndeliveredElementCatchingException(unwrapTyped(updCellResult))?.let { throw it }
                 }
             }
         }
@@ -865,7 +866,7 @@ internal open class BufferedChannel<E>(
         /* This lambda is invoked when an element has been
         successfully retrieved, either from the buffer or
         by making a rendezvous with a suspended sender. */
-        onElementRetrieved: (element: E) -> R,
+        onElementRetrieved: (element: Any?) -> R,
         /* This lambda is called when the operation suspends in the cell
         specified by the segment and its global and in-segment indices. */
         onSuspend: (segm: ChannelSegment<E>, i: Int, r: Long) -> R,
@@ -935,7 +936,7 @@ internal open class BufferedChannel<E>(
                     // Clean the reference to the previous segment before finishing.
                     segment.cleanPrev()
                     @Suppress("UNCHECKED_CAST")
-                    onElementRetrieved(updCellResult as E)
+                    onElementRetrieved(updCellResult)
                 }
             }
         }
@@ -953,7 +954,7 @@ internal open class BufferedChannel<E>(
         /* This lambda is invoked when an element has been
         successfully retrieved, either from the buffer or
         by making a rendezvous with a suspended sender. */
-        onElementRetrieved: (element: E) -> Unit,
+        onElementRetrieved: (element: Any?) -> Unit,
         /* This lambda is called when the channel is observed
         in the closed state and no waiting senders is found,
         which means that it is closed for receiving. */
@@ -1512,7 +1513,7 @@ internal open class BufferedChannel<E>(
     private fun registerSelectForReceive(select: SelectInstance<*>, ignoredParam: Any?) =
         receiveImpl( // <-- this is an inline function
             waiter = select,
-            onElementRetrieved = { elem -> select.selectInRegistrationPhase(elem) },
+            onElementRetrieved = { elem -> select.selectInRegistrationPhase(unwrapTyped(elem)) },
             onSuspend = { _, _, _ -> },
             onClosed = { onClosedSelectOnReceive(select) }
         )
@@ -1609,7 +1610,7 @@ internal open class BufferedChannel<E>(
                 // Also, inform the `BufferedChannel` extensions that
                 // the synchronization of this receive operation is completed.
                 onElementRetrieved = { element ->
-                    saveReceiveResult(element)
+                    this.receiveResult = element
                     true
                 },
                 // As no waiter is provided, suspension is impossible.
@@ -1647,9 +1648,9 @@ internal open class BufferedChannel<E>(
                 // In case `onUndeliveredElement` is present, we must
                 // invoke it in the latter case.
                 onElementRetrieved = { element ->
-                    saveReceiveResult(element)
+                    this.receiveResult = element
                     this.continuation = null
-                    cont.resume(true, onUndeliveredElement?.bindCancellationFun(element, cont.context))
+                    cont.resume(true, onUndeliveredElement?.bindCancellationFun(unwrapTyped(element), cont.context))
                 },
                 onClosed = { onClosedHasNextNoWaiterSuspend() }
             )
@@ -1704,17 +1705,13 @@ internal open class BufferedChannel<E>(
             val cont = this.continuation!!
             this.continuation = null
             // Store the retrieved element in `receiveResult`.
-            saveReceiveResult(element)
+            this.receiveResult = wrapInternal(element)
             // Try to resume this `hasNext()`. Importantly, the receiver coroutine
             // may be cancelled after it is successfully resumed but not dispatched yet.
             // In case `onUndeliveredElement` is specified, we need to invoke it in the latter case.
             return cont.tryResume0(true, onUndeliveredElement?.bindCancellationFun(element, cont.context))
         }
 
-        private fun saveReceiveResult(element: E) {
-            this.receiveResult = wrapInternal(element)
-        }
-
         fun tryResumeHasNextOnClosedChannel() {
             /*
              * Read the current continuation of the suspended `hasNext()` call and clean the corresponding field to avoid memory leaks.
@@ -2783,16 +2780,17 @@ internal class ChannelSegment<E>(id: Long, prev: ChannelSegment<E>?, channel: Bu
     }
 
     @Suppress("UNCHECKED_CAST")
-    internal fun getElement(index: Int) = data[index * 2].value as E
+    internal fun getElement(index: Int) = unwrapInternal(data[index * 2].value) as E
 
-    internal fun retrieveElement(index: Int): E = getElement(index).also { cleanElement(index) }
+    @Suppress("UNCHECKED_CAST")
+    internal fun retrieveElement(index: Int): E = (data[index * 2].value as E).also { cleanElement(index) }
 
     internal fun cleanElement(index: Int) {
-        setElementLazy(index, null)
+        data[index * 2].lazySet(null)
     }
 
     private fun setElementLazy(index: Int, value: Any?) {
-        data[index * 2].lazySet(value)
+        data[index * 2].lazySet(wrapInternal(value))
     }
 
     // ######################################

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

@@ -14,6 +14,7 @@ internal class FlowValueWrapperInternal<T>(val value: T)
 
 internal fun <T> wrapInternal(value: T): T = value
 internal fun <T> unwrapInternal(value: T): T = value
+internal fun <T> unwrapTyped(value: Any?): T = NULL.unbox(unwrapInternal(value))
 
 // debugger agent transforms wrapInternal so it returns wrapInternalDebuggerCapture(value) instead of just value.
 private fun wrapInternalDebuggerCapture(value: Any?): Any {
@@ -35,5 +36,5 @@ private fun unwrapInternalDebuggerCapture(value: Any?): Any? {
 
 // Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
 internal suspend fun <T> FlowCollector<T>.emitInternal(value: Any?) {
-    emit(NULL.unbox(unwrapInternal(value)))
+    emit(unwrapTyped<T>(value))
 }