In-place StateFlow implementations

InPlaceCombinedFlow.kt

import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.combine

/**
 * Replacement for [combine][kotlinx.coroutines.flow.combine] operation for [StateFlow]s.
 * Keeps the result as a [StateFlow], provides realtime results ([StateFlow] result).
 * Also reduces number of allocations.
 *
 * @see kotlinx.coroutines.flow.combine
 */
fun <T1, T2, R> StateFlow<T1>.combineInPlace(
    flow: StateFlow<T2>,
    @BuilderInference transform: (T1, T2) -> R,
): StateFlow<R> = InPlaceCombinedFlow(this, flow, transform)

/** WARN: The [StateFlow] interface is not long-term stable for inheritance. */
private class InPlaceCombinedFlow<T1, T2, R>(
    private val flow1: StateFlow<T1>,
    private val flow2: StateFlow<T2>,
    private val transform: (T1, T2) -> R,
) : StateFlow<R> {

    override val value get() = transform(flow1.value, flow2.value)


    /** @see kotlinx.coroutines.flow.StateFlowImpl.replayCache */
    override val replayCache: List<R> get() = listOf(value)

    override suspend fun collect(collector: FlowCollector<R>): Nothing {
        /**
         * @see kotlinx.coroutines.flow.combine
         * @see kotlinx.coroutines.flow.DistinctFlowImpl
         */
        var previous: Any? = NULL
        flow1.combine(flow2, transform).collect { newValue ->
            if (previous === NULL || previous != newValue) {
                previous = newValue
                collector.emit(newValue)
            }
        }
        error("Should never happen")
    }

    private companion object NULL
}

InPlaceMappedFlow.kt

import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.combine

/**
 * Replacement for [map][kotlinx.coroutines.flow.map] operation for [StateFlow]s.
 * Keeps the result as a [StateFlow], provides realtime results ([StateFlow] result).
 * Also reduces number of allocations.
 *
 * @see kotlinx.coroutines.flow.map
 */
fun <T, R> StateFlow<T>.mapInPlace(transform: (value: T) -> R): StateFlow<R> = InPlaceMappedFlow(this, transform)

/** WARN: The [StateFlow] interface is not long-term stable for inheritance. */
private class InPlaceMappedFlow<T, R>(
    private val upstream: StateFlow<T>,
    private val transform: (value: T) -> R,
) : StateFlow<R> {

    override val value get() = transform(upstream.value)

    /** @see kotlinx.coroutines.flow.StateFlowImpl.replayCache */
    override val replayCache: List<R> get() = listOf(value)

    override suspend fun collect(collector: FlowCollector<R>): Nothing {
        /**
         * @see kotlinx.coroutines.flow.map
         * @see kotlinx.coroutines.flow.DistinctFlowImpl
         */
        var previous: Any? = NULL
        upstream.collect { value ->
            val newValue = transform(value)
            if (previous === NULL || previous != newValue) {
                previous = newValue
                collector.emit(newValue)
            }
        }
    }

    private companion object NULL
}