Attempt at implementing RxJava's switchMap for Kotlin Coroutine Channels

ChannelSwitchMap.kt

import kotlinx.coroutines.experimental.Unconfined
import kotlinx.coroutines.experimental.channels.ReceiveChannel
import kotlinx.coroutines.experimental.channels.consumes
import kotlinx.coroutines.experimental.channels.produce
import kotlinx.coroutines.experimental.selects.whileSelect
import kotlin.coroutines.experimental.CoroutineContext

fun <T, R> ReceiveChannel<T>.switchMap(context: CoroutineContext = Unconfined,
                                       transform: suspend (T) -> ReceiveChannel<R>): ReceiveChannel<R> {
  val input = this

  return produce(context, onCompletion = consumes()) {
    var current: ReceiveChannel<R> = transform(input.receive())
    val output = this

    whileSelect {
      input.onReceiveOrNull { t: T? ->
        t?.also { current = transform(it) } != null
      }

      current.onReceiveOrNull { r ->
        r?.also { output.send(it) } != null
      }
    }
  }
}

ChannelSwitchMapTest.kt

import com.google.common.truth.Truth.assertThat
import kotlinx.coroutines.experimental.channels.SendChannel
import kotlinx.coroutines.experimental.channels.produce
import kotlinx.coroutines.experimental.channels.toList
import kotlinx.coroutines.experimental.delay
import kotlinx.coroutines.experimental.runBlocking
import org.junit.Test
import kotlin.coroutines.experimental.coroutineContext

suspend fun <E> SendChannel<E>.sendAndLog(name: String, e: E) {
  println("$name is sending $e on ${Thread.currentThread().name}")
  send(e)
}

class ChannelSwitchMapTest {
  @Test fun `it switchMaps`() = runBlocking<Unit> {
    val source = produce(coroutineContext) {
      send(1)
      delay(10)
      send(2)
      delay(100)
      send(3)
    }

    val switchMapped = source.switchMap(coroutineContext) { i ->
      produce {
        send("${i}A")
        send("${i}B")
        delay(15)
        send("${i}C")
      }
    }

    val expected = switchMapped.toList()
    assertThat(expected).isEqualTo(listOf("1A", "1B", "2A", "2B", "2C"))
  }
}

You should wait 3A, 3B, 3C and manage backpressure "on send".

My bugged version

https://gist.github.com/fvasco/09724edd594e0d0c41e8792fb69cbaad

@fvasco right, thanks!