What is the most concise way to include functions or lambdas as conditions in a Kotlin `when` statement (or other branching construct)?

Question

I'm processing strings, and I came across the Regex or Wildcard answer: that one can put regular expressions in a when statement with a custom class that overrides equals. While this does effectively use the type system to shoehorn syntactic sugar into the when statement, I find the following pretty ugly, and would never do this in code that I intend to share with another developer (quoting travis):

import kotlin.text.regex

when (RegexWhenArgument(uri)) {
    Regex(/* pattern */) -> /* do stuff */
    Regex(/* pattern */) -> /* do stuff */
    /* etc */
}

Where RegexWhenArgument is minimally defined as:

class RegexWhenArgument (val whenArgument: CharSequence) {
    operator fun equals(whenEntry: Regex) = whenEntry.matches(whenArgument)
    override operator fun equals(whenEntry: Any?) = (whenArgument == whenEntry)
}

(end quote)

I think it would be much more readable to pass an arg to when and then reference functions that operate on the type of the arg. For a contrived example:

// local declaration
val startsWithFn: (String) -> Boolean = {s -> s.startsWith("fn:")}

when(givenString) {
    ::startsWithHelp -> printHelp()
    startsWithFn -> println("Hello, ${givenString.substring(3)}!")
}

// package level function
fun startsWithHelp(s:String) = s.startsWith("help", true)

But of course, this code doesn't compile. Is there a way to do this that's readable, maintainable, and concise? Maybe using Streams? What would an experienced Kotlin developer do?

Answer 1

There are quite a few solutions to your problem. I will start with the simple ones and then move on to the more complex.

---

When without argument

Kotlin documentation says:

when can also be used as a replacement for an if-else if chain. If no argument is supplied, the branch conditions are simply boolean expressions, and a branch is executed when its condition is true

Use case:

when {
    startsWithHelp(givenString) -> printHelp()
    startsWithFn(givenString) -> println("Hello, ${givenString.substring(3)}!")
}

Pros:

1. Doesn't require any additional code

2. Easy to understand

Cons:

1. Boilerplate (givenString)

---

Argument wrapper which overrides equals

The required wrapper is similar to RegexWhenArgument, but instead of checking Regex it invokes (String) -> Boolean. There is nothing special about String, so I will use <T>.

Wrapper function:

fun <T> whenArg(arg: T) = object {
    override fun equals(other: Any?) =
        if (other is Function1<*, *>) 
            (other as Function1<T, Boolean>)(arg)
        else arg == other
}

Use case:

when (whenArg(givenString)) {
    ::startsWithHelp -> printHelp()
    startsWithFn -> println("Hello, ${givenString.substring(3)}!")
}

Pros:

1. Functions can be used in when branches
2. Very easy to understand
3. Requires only one simple wrapper function

Cons:

1. Easy to forget to wrap an argument
2. The function of the wrong type can be accidentally used in a branch condition

Note: this solution and the next one allows using functions combination, which can be created by these extensions:

infix fun <T> ((T) -> Boolean).and(other: ((T) -> Boolean)) = { it: T -> invoke(it) && other(it) }
infix fun <T> ((T) -> Boolean).or(other: ((T) -> Boolean)) = { it: T -> invoke(it) || other(it) }
operator fun <T> ((T) -> Boolean).not() = { it: T -> !invoke(it) }

---

DSL-based when analog

In Kotlin DSL can be used not only for building objects but for creating custom program flow structures as well.

Source code:

@DslMarker
annotation class WhichDsl

// This object is used for preventing client code from creating nested
// branches. You can omit it if you need them, but I highly recommend
// not to do this because nested branches may be confusing. 
@WhichDsl
object WhichCase

// R type parameter represents a type of expression result
@WhichDsl
class Which<T, R>(val arg: T) {
    // R? is not used here because R can be nullable itself 
    var result: Holder<R>? = null

    inline operator fun ((T) -> Boolean).invoke(code: WhichCase.() -> R) {
        if (result == null && invoke(arg)) result = Holder(code(WhichCase))
    }

    // else analog
    inline fun other(code: WhichCase.() -> R) = result?.element ?: code(WhichCase)
}

data class Holder<out T>(val element: T)

inline fun <T, R> which(arg: T, @BuilderInference code: Which<T, R>.() -> R) =
    Which<T, R>(arg).code()

Statement use case:

which(givenString) {
    ::startsWithHelp { printHelp() }
    startsWithFn { println("Hello, ${givenString.substring(3)}!") }
}

Expression use case:

val int = which(givenString) {
    ::startsWithHelp { 0 }
    startsWithFn { 1 }
    other { error("Unknown command: $givenString") }
}

Pros:

1. Customizable syntax
2. New features can be added
3. Functions which works only with some argument types can be added
4. Nested branches can be enabled
5. No boilerplate
6. Function types are checked at compile-time
7. Easy to read

Cons:

1. Requires a few helper classes and functions
2. Usage of experimental @BuilderInference (it can be replaced with explicit type declaration for expressions and separate method for statements)
3. New syntax learning may take some time

Extensions examples:

inline fun <R> Which<*, R>.orThrow(message: () -> String) =
    other { throw NoWhenBranchMatchedException(message()) }

val <R> Which<*, R>.orThrow get() = orThrow { "No branch matches to $arg" }

inline fun <T, R> Which<T, R>.branch(condition: (T) -> Boolean, code: WhichCase.() -> R) =
    condition(code)

inline fun <T, R> Which<T, R>.case(value: T, code: WhichCase.() -> R) =
    branch({ it == value }, code)

fun <T : CharSequence, R> Which<T, R>.regex(regex: Regex, code: WhichCase.() -> R) =
    branch({ regex.matches(it) }, code)

fun <T : Comparable<T>, R> Which<T, R>.range(range: ClosedRange<T>, code: WhichCase.() -> R) =
    branch({ it in range }, code)

inline fun <T> Which<T, *>.sideBranch(condition: (T) -> Boolean, code: WhichCase.() -> Unit) {
    if (condition(arg)) code(WhichCase)
}

fun <T> Which<T, *>.sideCase(value: T, code: WhichCase.() -> Unit) =
    sideBranch({ it == value }, code)

inline fun <R> Which<*, R>.dropResult(condition: WhichCase.(R) -> Boolean = { _ -> true }) {
    result?.let { (element) ->
        if (WhichCase.condition(element)) result = null
    }
}

inline fun <T, R> Which<T, R>.subWhich(condition: (T) -> Boolean, code: Which<T, R>.() -> R) =
    branch(condition) {
    which(this@subBranch.arg, code)
}