Support SelectImplementation

This brings support for select-based flow operators, such as `timeout`.

Author: AlexVanGogen

IntelliJ-patches.md

@@ -61,11 +61,9 @@ The agent needs three entities to establish a proper asynchronous stack traces c
 
 The key for MutableStateFlow is the element itself. For MutableSharedFlow, the element is wrapped into a unique object to prevent bridging mistakes when two equal elements are emitted from different places.
 
-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.
+Most of the operators applicable to flows (such as `map`, `scan`, `debounce`, `timeout`, `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 use `Select` inside (such as `timeout`).
-
 ### API
 
 Some logic related to instrumentation was extracted to separate methods so that the debugger agent could instrument it properly:
@@ -76,7 +74,8 @@ Some logic related to instrumentation was extracted to separate methods so that
 - `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
+- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.emitInternal(FlowCollector, value)` -- alternative of a regular `FlowCollector.emit` that supports insertion points; if there is a `FlowCollector`, its `emit` call can be replaced with `emitInternal` so this case would also be supported for constructing async stack traces 
+- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.debuggerCapture` -- common insertion point for a debugger agent; simplifies instrumentation; the value is always being unwrapped inside. `emitInternal` uses this method.
 
 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.
 
@@ -86,3 +85,5 @@ 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.
+
+`SelectImplementation.complete` now uses `debuggerCapture` to properly propagate value that might come from flows. 

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

@@ -1513,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(unwrapTyped(elem)) },
+            onElementRetrieved = { elem -> select.selectInRegistrationPhase(elem) },
             onSuspend = { _, _, _ -> },
             onClosed = { onClosedSelectOnReceive(select) }
         )

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

@@ -36,5 +36,10 @@ 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(unwrapTyped<T>(value))
+    debuggerCapture<T, Unit>(value) { emit(it) }
+}
+
+// Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
+internal suspend fun <Unwrapped, R> debuggerCapture(value: Any?, block: suspend (Unwrapped) -> R): R {
+    return block(unwrapTyped<Unwrapped>(value))
 }

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

@@ -229,7 +229,7 @@ private fun <T> Flow<T>.debounceInternal(timeoutMillisSelector: (T) -> Long): Fl
                 }
                 values.onReceiveCatching { value ->
                     value
-                        .onSuccess { lastValue = it }
+                        .onSuccess { lastValue = wrapInternal(it) }
                         .onFailure {
                             it?.let { throw it }
                             // If closed normally, emit the latest value
@@ -278,7 +278,7 @@ public fun <T> Flow<T>.sample(periodMillis: Long): Flow<T> {
             select<Unit> {
                 values.onReceiveCatching { result ->
                     result
-                        .onSuccess { lastValue = it }
+                        .onSuccess { lastValue = wrapInternal(it) }
                         .onFailure {
                             it?.let { throw it }
                             ticker.cancel(ChildCancelledException())

kotlinx-coroutines-core/common/src/selects/Select.kt

@@ -3,6 +3,8 @@ package kotlinx.coroutines.selects
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.internal.debuggerCapture
+import kotlinx.coroutines.flow.internal.wrapInternal
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.selects.TrySelectDetailedResult.*
 import kotlin.contracts.*
@@ -626,7 +628,7 @@ internal open class SelectImplementation<R>(
                         val cont = curState as CancellableContinuation<Unit>
                         // Success! Store the resumption value and
                         // try to resume the continuation.
-                        this.internalResult = internalResult
+                        this.internalResult = wrapInternal(internalResult)
                         if (cont.tryResume(onCancellation)) return TRY_SELECT_SUCCESSFUL
                         // If the resumption failed, we need to clean the [result] field to avoid memory leaks.
                         this.internalResult = NO_RESULT
@@ -690,19 +692,21 @@ internal open class SelectImplementation<R>(
         // of memory leaks. Collect the internal result before that.
         val internalResult = this.internalResult
         cleanup(selectedClause)
-        // Process the internal result and invoke the user's block.
-        return if (!RECOVER_STACK_TRACES) {
-            // TAIL-CALL OPTIMIZATION: the `suspend` block
-            // is invoked at the very end.
-            val blockArgument = selectedClause.processResult(internalResult)
-            selectedClause.invokeBlock(blockArgument)
-        } else {
-            // TAIL-CALL OPTIMIZATION: the `suspend`
-            // function is invoked at the very end.
-            // However, internally this `suspend` function
-            // constructs a state machine to recover a
-            // possible stack-trace.
-            processResultAndInvokeBlockRecoveringException(selectedClause, internalResult)
+        return debuggerCapture<Any?, R>(internalResult) { result ->
+            // Process the internal result and invoke the user's block.
+            if (!RECOVER_STACK_TRACES) {
+                // TAIL-CALL OPTIMIZATION: the `suspend` block
+                // is invoked at the very end.
+                val blockArgument = selectedClause.processResult(result)
+                selectedClause.invokeBlock(blockArgument)
+            } else {
+                // TAIL-CALL OPTIMIZATION: the `suspend`
+                // function is invoked at the very end.
+                // However, internally this `suspend` function
+                // constructs a state machine to recover a
+                // possible stack-trace.
+                processResultAndInvokeBlockRecoveringException(selectedClause, result)
+            }
         }
     }