Kotlinのforループで独自クラスをイテレートする方法

概要

Kotlinでは、規約に従った実装を行うことで言語の構文に独自のプログラムを柔軟に組み込むことができます。
この記事では、 rangeTo の規約によって、forループ内で独自クラスをいてレートする実装例をご紹介します。

rangeTo と ClosedRange

Kotlin の for文は以下のように記述されますが、以下の 1..3 の部分の正体は ClosedRange です。

for (i in 1..3) print("${i*i} ") // 1 4 9

Kotlinのコンパイラは start..end を start.rangeTo(end) という呼び出しに変換するので、 start の型に rangeTo 関数が適切に定義されていれば、独自のクラスを for文の中でイテレートすることができます。
適切な、とは以下のシグネチャーで、 ClosedRange を返却する関数です。

operator fun <T: Comparable<T>> T.rangeTo(that: T): ClosedRange<T>

実装例

以下の十二支を表す列挙型(ChineseZodiac)をfor文で使えるように実装します。

enum class ChineseZodiac(val char: String) {
    RAT("子"), OX("丑"), TIGER("寅"), RABBIT("卯"), DRAGON("辰"), SNAKE("巳"),
    HORSE("午"), SHEEP("未"), MONKEY("申"), COCK("酉"), DOG("戌"), BOAR("亥");    
}

まずは、実装コードの見通しをよくするために、12を法とした加算が行えるように以下の拡張関数を定義しておきます。

operator fun ChineseZodiac.plus(i: Int): ChineseZodiac {
    val ord = (this.ordinal + i) % 12
    return ChineseZodiac.values()[ord]
}

これも、規約に従った実装となっていて、 plus 関数を ChineseZodiac に追加することにより、+ による演算が可能となります。

println(RAT + 1) // OX

次に rangeTo を拡張関数として定義します。

operator fun ChineseZodiac.rangeTo(that: ChineseZodiac) = ChineseZodiacRange(this, that)

ChineseZodiacRange の定義は以下となります。
ChineseZodiac を型パラメータとして、 ClosedRange インタフェースを実装します。

class ChineseZodiacRange(override val start: ChineseZodiac,
                         override val endInclusive: ChineseZodiac): ClosedRange<ChineseZodiac>

forループで in演算子によって ClosedRange をイテレートするためには、実はもう一つ規約に沿う必要があります。それは iterator() メソッドで Iterator を返却することです。
ChineseZodiacRangeの拡張関数として以下のように実装します。

operator fun ChineseZodiacRange.iterator(): Iterator<ChineseZodiac> =
	object : Iterator<ChineseZodiac> {
        private var current = start
        override fun hasNext() = current != (endInclusive + 1)
        override fun next(): ChineseZodiac {
            val ret = current
            current = current + 1
            return ret
        }
    }

これで、以下のようにforループでイテレートできるようになりました。

    for (e in TIGER..DRAGON) print(e.char) // 寅卯辰
    println()
    for (e in COCK..OX) print(e.char) // 酉戌亥子丑
    println()
    for (e in OX..OX) print(e.char) // 丑
    println()

ソースコード

今回の実装例のコード一式です：

import ChineseZodiac.*

fun main(args: Array<String>) {
    for (e in TIGER..DRAGON) print(e.char) // 寅卯辰
    println()
    for (e in COCK..OX) print(e.char) // 酉戌亥子丑
    println()
    for (e in OX..OX) print(e.char) // 丑
    println()
}

enum class ChineseZodiac(val char: String) {
    RAT("子"), OX("丑"), TIGER("寅"), RABBIT("卯"), DRAGON("辰"), SNAKE("巳"),
    HORSE("午"), SHEEP("未"), MONKEY("申"), COCK("酉"), DOG("戌"), BOAR("亥");    
}

operator fun ChineseZodiac.plus(i: Int): ChineseZodiac {
    val ord = (this.ordinal + i) % 12
    return ChineseZodiac.values()[ord]
}

operator fun ChineseZodiac.rangeTo(that: ChineseZodiac) = ChineseZodiacRange(this, that)

class ChineseZodiacRange(override val start: ChineseZodiac,
                         override val endInclusive: ChineseZodiac): ClosedRange<ChineseZodiac>

operator fun ChineseZodiacRange.iterator(): Iterator<ChineseZodiac> =
	object : Iterator<ChineseZodiac> {
        private var current = start
        override fun hasNext() = current != (endInclusive + 1)
        override fun next(): ChineseZodiac {
            val ret = current
            current = current + 1
            return ret
        }
    }