monolithic kernel

key としてすべての enum を網羅した Map

プログラムを書いていると、enum class をキーとして、何らかの値を保持する Map を作ることがよくある。

enum class MyEnum {
    KEY1,
    KEY2,
}


val map = mapOf(
    MyEnum.KEY1 to "VALUE1",
    MyEnum.KEY2 to "VALUE2",
)

key としてすべての enum を網羅していた場合、Map#get が null を返すことはないはずなのだが、戻り値は nullable であるため、checkNotNull を挟むなどしなければならないのが面倒。

val v: String = checkNotNull(map[MyEnum.KEY1])

そこで、すべての enum を網羅していることを前提に、non-nullable な Map#get ができる仕組みを考えてみた。

import java.util.EnumMap


class CoveredEnumMap<K : Enum<K>, out V> @PublishedApi internal constructor(
    private val enumMap: EnumMap<K, V>,
) : Map<K, V> by enumMap {
    override operator fun get(key: K): V {
        return checkNotNull(enumMap.get(key))
    }


    override fun toString(): String {
        return "CoveredEnumMap($enumMap)"
    }


    override fun equals(other: Any?): Boolean {
        return enumMap.equals(other)
    }


    override fun hashCode(): Int {
        return enumMap.hashCode()
    }
}


inline fun <reified K : Enum<K>, V> coveredEnumMapOf(map: Map<K, V>): CoveredEnumMap<K, V> {
    require(map.size == enumValues<K>().size)
    return CoveredEnumMap(EnumMap(map))
}

基本的には enumMap に処理を移譲しつつ、Map#get の戻り値を non-nullable にしている。構築時にすべての enum を網羅していることを確認済みなので、Map#get したタイミングで例外が飛ぶことはない。

まあ、checkNotNull を wrap しただけといえばそうなのだが。

val map = coveredEnumMapOf(
    mapOf(
        MyEnum.KEY1 to "VALUE1",
        MyEnum.KEY2 to "VALUE2",
  )
)


val v: String = map[MyEnum.KEY1]
val map2: Map<MyEnum, String> = map

map を使う側では checkNotNull を使うことなく non-nullable な値を取得できる。Map の実装としてもそのまま使える。

これで Map#get については解決できたが、これだと enum を網羅していないことを実行時にしか検知できないのが惜しい。ぜひともコンパイル時に解決したいところ。

そこで考えたのが以下のコード。

inline fun <reified K : Enum<K>, V> coveredEnumMapOf(transform: (K) -> V): CoveredEnumMap<K, V> {
    return CoveredEnumMap(enumValues<K>().associateWithTo(EnumMap(K::class.java), transform))
}


val map = coveredEnumMapOf<MyEnum, String> {
    when (it) {
        MyEnum.KEY1 -> "VALUE1"
        MyEnum.KEY2 -> "VALUE2"
    }
}


val v: String = map[MyEnum.KEY1]

まず、各要素の key となる enum に対して値を決めさせるインターフェイスにすることで、強制的にすべての enum を key に持つ Map しか作れないようにした。値を決める部分は利用者次第だが、なにか共通の規則で導出するのでなければ、when 式が使える。when 式はすべての分岐を網羅していない場合にコンパイルエラーになるので、コンパイルが通っていれば記述に漏れがないことが確定し、実行時エラーの可能性を排除できる。

実用するかはともかくとして、イメージしたものを作れて満足した。

Related articles