Helper methods to build select clauses for Kotlin coroutines

SelectClauses.kt

package com.zachklipp.coroutines

import kotlinx.coroutines.experimental.DisposableHandle
import kotlinx.coroutines.experimental.intrinsics.startCoroutineCancellable
import kotlinx.coroutines.experimental.selects.SelectClause0
import kotlinx.coroutines.experimental.selects.SelectClause1
import kotlinx.coroutines.experimental.selects.SelectClause2
import kotlinx.coroutines.experimental.selects.SelectInstance

/**
 * Builds a clause that can be used in a select statement.
 *
 * Example:
 * ```
 * val Button.onClick: SelectClause0
 *   get() = buildNoArgSelectClause { builder ->
 *     val listener = View.OnClickListener { builder.trySelectMeAndResume() }
 *     builder.invokeOnSelection { setOnClickListener(null) }
 *     setOnClickListener(listener)
 *   }
 * ```
 *
 * @see buildSelectClauseWithOutput
 * @see buildSelectClauseWithInputAndOutput
 */
inline fun buildNoArgSelectClause(
  crossinline builderBlock: (NoArgSelectClauseBuilder<*>) -> Unit
): SelectClause0 = object : SelectClause0 {
  override fun <R> registerSelectClause0(
    select: SelectInstance<R>,
    block: suspend () -> R
  ) {
    // Short circuit if a previously-registered clause has synchronously selected itself.
    if (select.isSelected) return
    builderBlock(NoArgSelectClauseBuilder(select, block))
  }
}

/**
 * Builds a select clause that selects on a value of type `Q`.
 *
 * Example:
 * ```
 * val TextView.onTextChanged: SelectClause1
 *   get() = buildSelectClauseWithOutput { builder ->
 *     val listener = object : TextWatcher {
 *       override fun onTextChanged(s: CharSequence, start: Int, before: Int, count: Int) {
 *         builder.trySelectMeAndResume()
 *       }
 *       // Other methods are no-ops and omitted for brevity.
 *     }
 *     builder.invokeOnSelection { removeTextChangedListener(listener) }
 *     addTextChangedListener(listener)
 *   }
 * ```
 *
 * @see buildNoArgSelectClause
 * @see buildSelectClauseWithInputAndOutput
 */
inline fun <Q> buildSelectClauseWithOutput(
  crossinline builderBlock: (SelectClauseWithOutputBuilder<Q, *>) -> Unit
): SelectClause1<Q> = object : SelectClause1<Q> {
  override fun <R> registerSelectClause1(
    select: SelectInstance<R>,
    block: suspend (Q) -> R
  ) {
    // Short circuit if a previously-registered clause has synchronously selected itself.
    if (select.isSelected) return
    builderBlock(SelectClauseWithOutputBuilder(select, block))
  }
}

/**
 * Builds a select clause that operates on a value of type `P` and selects on a value of type `Q`.
 *
 * @see buildNoArgSelectClause
 * @see buildSelectClauseWithOutput
 */
inline fun <P, Q> buildSelectClauseWithInputAndOutput(
  crossinline builderBlock: (param: P, SelectClauseWithOutputBuilder<Q, *>) -> Unit
): SelectClause2<P, Q> = object : SelectClause2<P, Q> {
  override fun <R> registerSelectClause2(
    select: SelectInstance<R>,
    param: P,
    block: suspend (Q) -> R
  ) {
    // Short circuit if a previously-registered clause has synchronously selected itself.
    if (select.isSelected) return
    builderBlock(param, SelectClauseWithOutputBuilder(select, block))
  }
}

class NoArgSelectClauseBuilder<R>(
  @PublishedApi internal val select: SelectInstance<R>,
  @PublishedApi internal val block: suspend () -> R
) {
  /**
   * Attempts to mark this clause as the selected one, and if successful, resume the continuation.
   *
   * @param idempotenceKey An arbitrary value that will be used to ensure idempotence somehow – not
   * entirely sure exactly how this is used or in what case there could be a race.
   * @param onThisClauseSelected Invoked if this clause won the race, before resuming the
   * continuation.
   */
  inline fun trySelectMeAndResume(
    idempotenceKey: Any? = null,
    onThisClauseSelected: () -> Unit = {}
  ) {
    if (select.trySelect(idempotenceKey)) {
      onThisClauseSelected()
      block.startCoroutineCancellable(select.completion)
    }
  }

  /**
   * Registers a block to run when _any_ clause is selected and the continuation can be resumed.
   * `block` is invoked whether or not this particular clause was the one that got selected.
   */
  inline fun invokeOnSelection(crossinline block: () -> Unit) {
    select.disposeOnSelect(object : DisposableHandle {
      override fun dispose() {
        block()
      }
    })
  }
}

class SelectClauseWithOutputBuilder<in Q, R>(
  @PublishedApi internal val select: SelectInstance<R>,
  @PublishedApi internal val block: suspend (Q) -> R
) {
  /**
   * Attempts to mark this clause as the selected one, and if successful, resume the continuation.
   *
   * @param value The selected value that will be passed to the block associated with this clause.
   * @param idempotenceKey An arbitrary value that will be used to ensure idempotence somehow – not
   * entirely sure exactly how this is used or in what case there could be a race.
   * @param onThisClauseSelected Invoked if this clause won the race, before resuming the
   * continuation.
   */
  inline fun trySelectMeAndResume(
    value: Q,
    idempotenceKey: Any? = null,
    onThisClauseSelected: () -> Unit = {}
  ) {
    if (select.trySelect(idempotenceKey)) {
      onThisClauseSelected()
      block.startCoroutineCancellable(value, select.completion)
    }
  }

  /**
   * Registers a block to run when _any_ clause is selected and the continuation can be resumed.
   * `block` is invoked whether or not this particular clause was the one that got selected.
   */
  inline fun invokeOnSelection(crossinline block: () -> Unit) {
    select.disposeOnSelect(object : DisposableHandle {
      override fun dispose() {
        block()
      }
    })
  }
}