【Kotlin/JVM】enumの複数判定を少し簡潔に書けたらいいな

（2020-03-08 16:00 訂正あり）

「enumの複数判定を少し簡潔に書けたらいいな」@JmzSpR の Kotlin/JVM 版。

Java の EnumSet を使っているため Kotlin/JVM 限定だが、
EnumSet は実装がそれほど難しくないので、自作すれば他の環境でも同様にできるだろう。

課題

ある enum が定義されているとき、その型のある値が、その型の値のある集合に属しているかどうかをテストしたい。

たとえばつぎのような enum が定義されているとして

enum class Color {
    RED, GREEN, BLUE
}

ある変数 val color: Color の値が RED もしくは BLUE であるかどうかをテストしたい。

簡潔に書けるようにするにはどうしたらよいか。

回答案

Enum の + 演算子をオーバーロードする

次のような拡張関数を実装する。

/**
 * [this] と [enum] だけを含む [EnumSet] を返す。
 */
inline operator fun <reified T : Enum<T>> T.plus(enum: T): EnumSet<T> =
    EnumSet.of(this, enum)

これにより、RED + BLUE という演算で RED と BLUE だけを含む EnumSet<Color> インスタンスが返るようになる。

これで次のように書ける。

fun main() {
    if (RED in RED + BLUE) {
        println("$RED は ${RED + BLUE} に含まれています。")
    } else {
        throw AssertionError()
    }

    if (RED in RED + BLUE + GREEN - RED) {
        throw AssertionError()
    } else {
        println("$RED は　${RED + BLUE + GREEN - RED} に含まれていません。")
    }
}

実行結果の標準出力：

RED は [RED, BLUE] に含まれています。
RED は　[BLUE, GREEN] に含まれていません。

なお、EnumSet に対する - 演算は Set インターフェイスの拡張関数として、in 演算は Collection インターフェイスでの contains メンバ関数のオーバーロードとして、標準で提供されている。

Enum の containts メンバ関数（in 演算子）をオーバーロードする

前項の変更だけだと、RED in RED のような enum 単体に対しての場合に使えない。
これも同じようにあつかえるようにするため、次の実装を追加する。

/**
 * [this] と [enum] が等価であるかどうかを返す。
 */
operator fun <T : Enum<T>> T.contains(enum: T): Boolean =
    this == enum

これで次のように書ける。

fun main() {
    if (RED in RED) {
        println("$RED は $RED に含まれています。")
    } else {
        throw AssertionError()
    }

    if (RED in BLUE) {
        throw AssertionError()
    } else {
        println("$RED は $BLUE に含まれていません。")
    }
}

実行結果の標準出力：

RED は RED に含まれています。
RED は BLUE に含まれていません。

Enum の - 演算子をオーバーロードする

RED in RED - BLUE のように、+ 演算によって EnumSet になる前に - 演算が行われた場合に対応できるようにする。

/**
 * [this] と [enum] が等価であれば空の、そうでなければ [this] だけを含む [EnumSet] を返す。
 */
inline operator fun <reified T : Enum<T>> T.minus(enum: T): EnumSet<T> =
    if (this == enum) EnumSet.noneOf(T::class.java)
    else EnumSet.of(this)

これで次のように書ける。

fun main() {
    if (RED in RED - BLUE) {
        println("$RED は ${RED - BLUE} に含まれています。")
    } else {
        throw  AssertionError()
    }

    if (RED in RED - RED) {
        throw AssertionError()
    } else {
        println("$RED は ${RED - RED} に含まれていません。")
    }
}

実行結果の標準出力：

RED は [RED] に含まれています。
RED は [] に含まれていません。

（2020-03-08 00:32 追記）EnumSet の + - 演算子をオーバーロードする

ここまでの実装でも正しく動作するが、EnumSet に対して + 演算子や - 演算子を呼び出されると、Set の拡張関数が呼び出され、返値が EnumSet でない Set になってしまう。EnumSet は enum を扱うのに特化していて高速・省メモリなので、これは避けたい。

そこで次のようなコードを追加する。

/**
 * [this] に [enum] を加えた新しい [EnumSet] を返す。
 */
operator fun <T : Enum<T>> EnumSet<T>.plus(enum: T): EnumSet<T> =
    clone().also {
        it += enum
    }

/**
 * [this] から [enum] を除いた新しい [EnumSet] を返す。
 */
operator fun <T : Enum<T>> EnumSet<T>.minus(enum: T): EnumSet<T> =
    clone().also {
        it -= enum
    }

欠点

この方法だと、簡潔に書けるが、次のような欠点がある。

• enum に対する + - 演算の返値型が EnumSet 型である。
  • 通常、+ - 演算の返値型は非演算子と同じであるため、期待に反する動作となる。
• + - 演算をするたびに新しい EnumSet インスタンスを生成する。
  • （2020-03-08 16:00 訂正）EnumSet は enum の要素数と同じ長さの配列を内部に持つ。ほとんどの場合は問題にならないだろうが、要素数が非常に多い enum 型の場合や、パフォーマンスが重要な処理では考慮が必要となるだろう。（これは勘違い、EnumMap の話だ。）
    EnumSet はごく軽量なのでまず問題にならない。

まとめ

次のような定義をしておくことで

/**
 * [this] と [enum] だけを含む [EnumSet] を返す。
 */
inline operator fun <reified T : Enum<T>> T.plus(enum: T): EnumSet<T> =
    EnumSet.of(this, enum)

/**
 * [this] と [enum] が等価であれば空の、そうでなければ [this] だけを含む [EnumSet] を返す。
 */
inline operator fun <reified T : Enum<T>> T.minus(enum: T): EnumSet<T> =
    if (this == enum) EnumSet.noneOf(T::class.java)
    else EnumSet.of(this)

/**
 * [this] と [enum] が等価であるかどうかを返す。
 */
operator fun <T : Enum<T>> T.contains(enum: T): Boolean =
    this == enum

/**
 * [this] に [enum] を加えた新しい [EnumSet] を返す。
 */
operator fun <T : Enum<T>> EnumSet<T>.plus(enum: T): EnumSet<T> =
    clone().also {
        it += enum
    }

/**
 * [this] から [enum] を除いた新しい [EnumSet] を返す。
 */
operator fun <T : Enum<T>> EnumSet<T>.minus(enum: T): EnumSet<T> =
    clone().also {
        it -= enum
    }

RED in RED + GREEN - BLUE のような演算を行うことができる。