Repository: Kotlin/kotlinx.coroutines
File: kotlinx-coroutines-core/jvm/src/flow/internal/SafeCollector.kt (Kotlin)

Code before:
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*
import kotlin.coroutines.*
import kotlin.coroutines.intrinsics.*
import kotlin.coroutines.jvm.internal.*

@Suppress("UNCHECKED_CAST")
private val emitFun =
    FlowCollector<Any?>::emit as Function3<FlowCollector<Any?>, Any?, Continuation<Unit>, Any?>

/**
 * A safe collector is an instance of [FlowCollector] that ensures that neither context preservation
 * nor exception transparency invariants are broken. Instances of [SafeCollector] are used in flow
 * operators that provide raw access to the [FlowCollector] e.g. [Flow.transform].
 * Mechanically, each [emit] call captures [currentCoroutineContext], ensures it is not different from the
 * previously caught one and proceeds further. If an exception is thrown from the downstream,
 * it is caught, and any further attempts to [emit] lead to the [IllegalStateException].
 *
 * ### Performance hacks
 *
 * Implementor of [ContinuationImpl] (that will be preserved as ABI nearly forever)
 * in order to properly control `intercepted()` lifecycle.
 * The safe collector implements [ContinuationImpl] to pretend it *is* a state-machine of its own `emit` method.
 * It is [ContinuationImpl] and not any other [Continuation] subclass because only [ContinuationImpl] supports `intercepted()` caching.
 * This is the most performance-sensitive place in the overall flow pipeline, because otherwise safe collector is forced to allocate
 * a state machine on each element being emitted, for an each intermediate stage where safe collector is present.
 *
 * See a comment to [emit] for the explanation of what and how is being optimized.
 */
@Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER", "INVISIBLE_MEMBER", "INVISIBLE_REFERENCE", "UNCHECKED_CAST")
internal actual class SafeCollector<T> actual constructor(
    @JvmField internal actual val collector: FlowCollector<T>,
    @JvmField internal actual val collectContext: CoroutineContext
) : FlowCollector<T>, ContinuationImpl(NoOpContinuation, EmptyCoroutineContext), CoroutineStackFrame {

    override val callerFrame: CoroutineStackFrame? get() = completion as? CoroutineStackFrame

    override fun getStackTraceElement(): StackTraceElement? = null

    @JvmField // Note, it is non-capturing lambda, so no extra allocation during init of SafeCollector
    internal actual val collectContextSize = collectContext.fold(0) { count, _ -> count + 1 }

    // Either context of the last emission or wrapper 'DownstreamExceptionContext'
    private var lastEmissionContext: CoroutineContext? = null
    // Completion if we are currently suspended or within completion body or null otherwise
    private var completion: Continuation<Unit>? = null

    /*
     * This property is accessed in two places:
     * * ContinuationImpl invokes this in its `releaseIntercepted` as `context[ContinuationInterceptor]!!`
     * * When we are within a callee, it is used to create its continuation object with this collector as completion

Review comment (suggestion):
```suggestion
 * a state machine on each element being emitted for each intermediate stage where the safe collector is present.
```

Code after:
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*
import kotlin.coroutines.*
import kotlin.coroutines.intrinsics.*
import kotlin.coroutines.jvm.internal.*

@Suppress("UNCHECKED_CAST")
private val emitFun =
    FlowCollector<Any?>::emit as Function3<FlowCollector<Any?>, Any?, Continuation<Unit>, Any?>

/**
 * A safe collector is an instance of [FlowCollector] that ensures that neither context preservation
 * nor exception transparency invariants are broken. Instances of [SafeCollector] are used in flow
 * operators that provide raw access to the [FlowCollector] e.g. [Flow.transform].
 * Mechanically, each [emit] call captures [currentCoroutineContext], ensures it is not different from the
 * previously caught one and proceeds further. If an exception is thrown from the downstream,
 * it is caught, and any further attempts to [emit] lead to the [IllegalStateException].
 *
 * ### Performance hacks
 *
 * Implementor of [ContinuationImpl] (that will be preserved as ABI nearly forever)
 * in order to properly control `intercepted()` lifecycle.
 * The safe collector implements [ContinuationImpl] to pretend it *is* a state-machine of its own `emit` method.
 * It is [ContinuationImpl] and not any other [Continuation] subclass because only [ContinuationImpl] supports `intercepted()` caching.
 * This is the most performance-sensitive place in the overall flow pipeline, because otherwise safe collector is forced to allocate
 * a state machine on each element being emitted for each intermediate stage where the safe collector is present.
 *
 * See a comment to [emit] for the explanation of what and how is being optimized.
 */
@Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER", "INVISIBLE_MEMBER", "INVISIBLE_REFERENCE", "UNCHECKED_CAST")
internal actual class SafeCollector<T> actual constructor(
    @JvmField internal actual val collector: FlowCollector<T>,
    @JvmField internal actual val collectContext: CoroutineContext
) : FlowCollector<T>, ContinuationImpl(NoOpContinuation, EmptyCoroutineContext), CoroutineStackFrame {

    override val callerFrame: CoroutineStackFrame? get() = completion_ as? CoroutineStackFrame

    override fun getStackTraceElement(): StackTraceElement? = null

    @JvmField // Note, it is non-capturing lambda, so no extra allocation during init of SafeCollector
    internal actual val collectContextSize = collectContext.fold(0) { count, _ -> count + 1 }

    // Either context of the last emission or wrapper 'DownstreamExceptionContext'
    private var lastEmissionContext: CoroutineContext? = null
    // Completion if we are currently suspended or within completion_ body or null otherwise
    private var completion_: Continuation<Unit>? = null

    /*
     * This property is accessed in two places:
     * * ContinuationImpl invokes this in its `releaseIntercepted` as `context[ContinuationInterceptor]!!`
     * * When we are within a callee, it is used to create its continuation object with this collector as completion_