KotlinでPartialFunctionを実装しcollect, collectFirstしたい

Scalaでいうcollect的な関数がほしいことがあって、このような関数を実装してみた。

fun <A, B> List<A>.collect(vararg patterns: Pair<(A) -> Boolean, (A) -> B>): List<B> = this.mapNotNull { a ->
    patterns.firstOrNull { (condition, _) -> condition(a) }?.let { (_, transform) -> transform(a) }
}

使ってみるとこんな感じ。
ぱっと見は悪くはないが、 Pair({ it > 0 }, { "Positive: $it" }) を { it > 0 } to { "Positive: $it" } と書こうとするとコンパイルエラーになったり、条件が増えたり変換が複雑になるとテストがしづらくfatになり、保守性の低下につながるという問題がある。

val list = listOf(-1, 0, 1, 2, -3, 4)

val positiveOrZero = list.collect(
    Pair({ it > 0 }, { "Positive: $it" }),
    Pair({ it == 0 }, { "Zero" })
)

println(positiveOrZero)
// ["Zero", "Positive: 1", "Positive: 2", "Positive: 4"]

ゆえにやはりPartialFunctionを引数に取るのが良いだろうということになり、KotlinにはPartialFunctionがないので実装してみることにする。

PartialFunctionの実装

PartialFunctionつまり部分関数の説明はこのへんを参考に。

関数の中でも部分関数というのが出てきますが。この定義をわかりや... - Yahoo!知恵袋

部分関数 - mrsekut-p

collectを実装するのに最低限必要な関数だけを持ったPartialFunctionがこちら。
compositeやandThenなどはないが、collectに限らず大抵のPartialFunctionの要件には充分なはず。
ポイントは (A) -> B? を継承しておりinvokeの戻りがnullableである点で、これにより PartialFunctionの "すべてのA型の値に対してB型の値に変換できる必要はない" という性質を表現している。

interface PartialFunction<A, B> : (A) -> B? {
    fun isDefinedAt(a: A): Boolean
    fun apply(a: A): B
    override fun invoke(a: A): B? = if (isDefinedAt(a)) apply(a) else null

    fun orElse(other: PartialFunction<A, B>): PartialFunction<A, B> = object : PartialFunction<A, B> {
        override fun isDefinedAt(a: A): Boolean = this@PartialFunction.isDefinedAt(a) || other.isDefinedAt(a)
        override fun apply(a: A): B = when {
            this@PartialFunction.isDefinedAt(a) -> this@PartialFunction.apply(a)
            else -> other.apply(a)
        }
    }
}

たとえばInt型の値に対して偶数の場合に文字列に変換するPartialFunctionは以下のように宣言できる。

val even = object : PartialFunction<Int, String> {
    override fun isDefinedAt(a: Int): Boolean = a % 2 == 0
    override fun apply(a: Int): String = "Even: $a"
}

ただ、PartialFunctionの生成を行うのに無名クラスの宣言を行う必要があり冗長になるので、以下のようなPartialFunction生成用関数partialFunctionやその省略表記pfを用意しておく。
これとPartialFunction.orElseを組み合わせることで条件分岐のような書き方になり、前述のcollectの例のようにできて見通しも良くなるはずである。

fun <A, B> partialFunction(isDefinedAt: (A) -> Boolean, apply: (A) -> B): PF<A, B> =
    object : PartialFunction<A, B> {
        override fun isDefinedAt(a: A): Boolean = isDefinedAt(a)
        override fun apply(a: A) = apply(a)
    }

typealias PF<A, B> = PartialFunction<A, B>

fun <A, B> pf(isDefinedAt: (A) -> Boolean, apply: (A) -> B): PF<A, B> =
    partialFunction(isDefinedAt, apply)

PartialFunctionを利用したcollect関数

PartialFunctionをinvokeした結果として条件を満たさない場合はnullを返すので、単にmapNotNullの内側でPartialFunctionを呼び出してあげるだけでcollectが出来上がる。

fun <A, B> List<A>.collect(pf: PartialFunction<A, B>): List<B> =
    this.mapNotNull { pf(it) }

なお、mapNotNullをfirstNotNullOfOrNullにしてあげるとcollectFirstとなる。

fun <A, B> List<A>.collectFirst(pf: PartialFunction<A, B>): B? = 
    this.firstNotNullOfOrNull { pf(it) }

collectの利用はこのように。

val list = listOf(-1, 0, 1, 2, -3, 4)

val positive: PF<Int, String> = pf({ it > 0 }) { "Positive: $it" }
val zero: PF<Int, String> = pf({ it == 0 }) { "Zero" }
val positiveOrZero = list.collect(positive.orElse(zero))

println(positiveOrZero)
// ["Zero", "Positive: 1", "Positive: 2", "Positive: 4"]

PartialFunctionであるpositive, zeroを生成してorElseでチェーンさせることで条件分岐を表現できている。
また、positive, zeroが独立した関数になっているため、それぞれの関数を独立して小さくテスト可能ということでありテスト容易性と関数の変更容易性が高くなる。この程度ならあまり気にならないが、パターンが複雑になるとメリットが感じられるはず。

その他PartialFunctionの使いどころ

Resultが特定のエラーの場合のみ回収する関数なども簡単に実装できて良い。
これもScalaのTry[T].recoverの挙動に近くなる。

fun <T> Result<T>.recover(pf: PartialFunction<Throwable, T>): Result<T> =
    when (val exception = exceptionOrNull()) {
        null -> this
        else -> when (val recovered = pf(exception)) {
            null -> this
            else -> result { recovered }
        }
    }