ReceiveChannel<E>.switchMap

switchMap.kt

fun <E : Any, R> ReceiveChannel<E>.switchMap(context: CoroutineContext = DefaultDispatcher,
                                             transform: suspend (E) -> ReceiveChannel<R>): ReceiveChannel<R> =
        produce(context) {
            var originChannel: ReceiveChannel<E>? = this@switchMap
            var itemChannel = this@switchMap.receiveOrNull()?.let { transform(it) } ?: return@produce

            var itemToSend: R? = null
            whileSelect {
                originChannel?.onReceiveOrNull?.invoke { newItem ->
                    if (newItem == null) originChannel = null // origin closed
                    else itemChannel = transform(newItem)// receive on new channel
                    itemToSend = null
                    true
                }

                when {
                    itemToSend != null -> {
                        this@produce.onSend(itemToSend as R) {
                            itemToSend = null
                            true
                        }
                    }
                    else -> {
                        itemChannel.onReceiveOrNull { item ->
                            if (item == null) itemChannel = originChannel?.receiveOrNull()?.let { transform(it) } ?: return@onReceiveOrNull false
                            else itemToSend = item
                            true
                        }
                    }
                }
            }
        }