Kotlin Coroutines zip three Flows

Viewed 4038

There is zip function to zip two Flows. Is there something to zip three (or more) Flows together?

If not, can you help me to implement extension function for it? Something like:

flow.zip(flow2, flow3) { a, b, c -> 

}
4 Answers

I am new to Flows but this seems to be working for me.

    // This will hold all 3 values
    data class Foo(val i: Int, val j: Int, val k: Int)

    val flow1 = (1..10).asFlow()
    val flow2 = (11..20).asFlow()
    val flow3 = (21..30).asFlow()

    val combinedFlow = flow1.zip(flow2) {i, j ->  
        Pair(i, j)
    }.zip(flow3) {pair, k ->
        Foo(pair.first, pair.second, k)
    }

Then you would collect them and get the values out like this:

        viewModelScope.launch {
            repo.combinedFlow.collect {foo ->
                System.out.println(foo.i)
                System.out.println(foo.j)
                System.out.println(foo.k)
            }
        }

I have not tested this, but you can give it a try. There's a lot of underlying code for zip, so to exploit that I'm zipping the first two Flows into a Flow of Pairs, and then zip the Flow of Pairs to the third Flow. But the lambda passed to this function gets the first two already separated so it doesn't have to know about the intermediate Pair step.

fun <T1, T2, T3, R> zip(
    first: Flow<T1>,
    second: Flow<T2>,
    third: Flow<T3>,
    transform: suspend (T1, T2, T3) -> R
): Flow<R> =
    first.zip(second) { a, b -> a to b }
        .zip(third) { (a, b), c ->
            transform(a, b, c)
        }

Usage like this:

zip(flow1, flow2, flow3) { a, b, c ->
    Triple(a, b, c)
}

And here's an untested version for an arbitrary number of flows, but they have to be of the same type:

fun <T, R> zip(
    vararg flows: Flow<T>,
    transform: suspend (List<T>) -> R
): Flow<R> = when(flows.size) {
    0 -> error("No flows")
    1 -> flows[0].map{ transform(listOf(it)) }
    2 -> flows[0].zip(flows[1]) { a, b -> transform(listOf(a, b)) }
    else -> {
        var accFlow: Flow<List<T>> = flows[0].zip(flows[1]) { a, b -> listOf(a, b) }
        for (i in 2 until flows.size) {
            accFlow = accFlow.zip(flows[i]) { list, it ->
                list + it
            }
        }
        accFlow.map(transform)
    }
}

You can check the zip operator implementation and try to copy/emulate how it works adapting it to your needs.

Test it and make all the changes you need

fun <T1, T2, T3, R> Flow<T1>.zip(flow2: Flow<T2>, flow3: Flow<T3>, transform: suspend (T1, T2, T3) -> R): Flow<R> = channelFlow {

    val first: ReceiveChannel<T1> = produce {
        this@zip.collect {
            channel.send(it)
        }
    }

    val second: ReceiveChannel<T2> = produce {
        flow2.collect {
            channel.send(it)
        }
    }

    val third: ReceiveChannel<T3> = produce {
        flow3.collect {
            channel.send(it)
        }
    }

    (second as SendChannel<*>).invokeOnClose {
        if (!first.isClosedForReceive) first.cancel(MyFlowException())
        if (!third.isClosedForReceive) third.cancel(MyFlowException())
    }

    (third as SendChannel<*>).invokeOnClose {
        if (!first.isClosedForReceive) first.cancel(MyFlowException())
        if (!second.isClosedForReceive) second.cancel(MyFlowException())
    }

    val otherIterator = second.iterator()
    val anotherIterator = third.iterator()

    try {
        first.consumeEach { value ->
            if (!otherIterator.hasNext() || !anotherIterator.hasNext()) {
                return@consumeEach
            }
            send(transform(value, otherIterator.next(), anotherIterator.next()))
        }
    } catch (e: MyFlowException) {
        // complete
    } finally {
        if (!second.isClosedForReceive) second.cancel(MyFlowException())
        if (!third.isClosedForReceive) third.cancel(MyFlowException())
    }
}

class MyFlowException: CancellationException()

Usage:

flow1.zip(flow2, flow3) { a, b, c ->
    //Do your work
}...

Flow has combine function

fun <T1, T2, R> Flow<T1>.combine(
    flow: Flow<T2>,
    transform: suspend (a: T1, b: T2) -> R
): Flow<R>

It returns a Flow whose values are generated with transform function by combining the most recently emitted values by each flow. Source

Example from srouces:

 val flow = flowOf(1, 2).onEach { delay(10) }
 val flow2 = flowOf("a", "b", "c").onEach { delay(15) }
 combine(flow, flow2) { i, s -> i.toString() + s }.collect {
     println(it) // Will print "1a 2a 2b 2c"
 }
Related