commonMain.flow.internal.Combine.kt Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of kotlinx-coroutines-core-common Show documentation
Coroutines support libraries for Kotlin
There is a newer version: 1.3.8
Show newest version
/*
 * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
 */
@file:Suppress("UNCHECKED_CAST", "NON_APPLICABLE_CALL_FOR_BUILDER_INFERENCE") // KT-32203

package kotlinx.coroutines.flow.internal

import kotlinx.coroutines.*
import kotlinx.coroutines.channels.*
import kotlinx.coroutines.flow.*
import kotlinx.coroutines.internal.*
import kotlinx.coroutines.selects.*

internal fun getNull(): Symbol = NULL // Workaround for JS BE bug

internal suspend fun  FlowCollector.combineTransformInternal(
    first: Flow, second: Flow,
    transform: suspend FlowCollector.(a: T1, b: T2) -> Unit
) {
    coroutineScope {
        val firstChannel = asFairChannel(first)
        val secondChannel = asFairChannel(second)
        var firstValue: Any? = null
        var secondValue: Any? = null
        var firstIsClosed = false
        var secondIsClosed = false
        while (!firstIsClosed || !secondIsClosed) {
            select {
                onReceive(firstIsClosed, firstChannel, { firstIsClosed = true }) { value ->
                    firstValue = value
                    if (secondValue !== null) {
                        transform(getNull().unbox(firstValue), getNull().unbox(secondValue) as T2)
                    }
                }

                onReceive(secondIsClosed, secondChannel, { secondIsClosed = true }) { value ->
                    secondValue = value
                    if (firstValue !== null) {
                        transform(getNull().unbox(firstValue) as T1, getNull().unbox(secondValue) as T2)
                    }
                }
            }
        }
    }
}

@PublishedApi
internal suspend fun  FlowCollector.combineInternal(
    flows: Array>,
    arrayFactory: () -> Array,
    transform: suspend FlowCollector.(Array) -> Unit
) {
    coroutineScope {
        val size = flows.size
        val channels =
            Array(size) { asFairChannel(flows[it]) }
        val latestValues = arrayOfNulls(size)
        val isClosed = Array(size) { false }

        // See flow.combine(other) for explanation.
        while (!isClosed.all { it }) {
            select {
                for (i in 0 until size) {
                    onReceive(isClosed[i], channels[i], { isClosed[i] = true }) { value ->
                        latestValues[i] = value
                        if (latestValues.all { it !== null }) {
                            val arguments = arrayFactory()
                            for (index in 0 until size) {
                                arguments[index] = NULL.unbox(latestValues[index])
                            }
                            transform(arguments as Array)
                        }
                    }
                }
            }
        }
    }
}

private inline fun SelectBuilder.onReceive(
    isClosed: Boolean,
    channel: ReceiveChannel,
    crossinline onClosed: () -> Unit,
    noinline onReceive: suspend (value: Any) -> Unit
) {
    if (isClosed) return
    channel.onReceiveOrNull {
        // TODO onReceiveOrClosed when boxing issues are fixed
        if (it === null) onClosed()
        else onReceive(it)
    }
}

// Channel has any type due to onReceiveOrNull. This will be fixed after receiveOrClosed
private fun CoroutineScope.asFairChannel(flow: Flow<*>): ReceiveChannel = produce {
    val channel = channel as ChannelCoroutine
    flow.collect { value ->
        return@collect channel.sendFair(value ?: NULL)
    }
}

internal fun  zipImpl(flow: Flow, flow2: Flow, transform: suspend (T1, T2) -> R): Flow = unsafeFlow {
    coroutineScope {
        val first = asChannel(flow)
        val second = asChannel(flow2)
        /*
         * This approach only works with rendezvous channel and is required to enforce correctness
         * in the following scenario:
         * ```
         * val f1 = flow { emit(1); delay(Long.MAX_VALUE) }
         * val f2 = flowOf(1)
         * f1.zip(f2) { ... }
         * ```
         *
         * Invariant: this clause is invoked only when all elements from the channel were processed (=> rendezvous restriction).
         */
        (second as SendChannel<*>).invokeOnClose {
            if (!first.isClosedForReceive) first.cancel(AbortFlowException())
        }

        val otherIterator = second.iterator()
        try {
            first.consumeEach { value ->
                if (!otherIterator.hasNext()) {
                    return@consumeEach
                }
                emit(transform(NULL.unbox(value), NULL.unbox(otherIterator.next())))
            }
        } catch (e: AbortFlowException) {
            // complete
        } finally {
            if (!second.isClosedForReceive) second.cancel(AbortFlowException())
        }
    }
}

// Channel has any type due to onReceiveOrNull. This will be fixed after receiveOrClosed
private fun CoroutineScope.asChannel(flow: Flow<*>): ReceiveChannel = produce {
    flow.collect { value ->
        return@collect channel.send(value ?: NULL)
    }
}