Creating a Kotlin Range Iterator on a Custom Object

1. Introduction

In a previous article, we showed how we can create a range in Kotlin, and how easy it is to iterate over Int, Long and Char types.

But what if we want to iterate over a custom type? Is it possible? The answer is yes! So let’s jump into the code and see how.

2. A Colorful Type

Let’s imagine we have a simple class that represents an RGB color:

class CustomColor(val rgb: Int): Comparable<CustomColor> {}

It’d be nice to be able to iterate over a range of RGB colors:

val a = CustomColor(0x000000)
val b = CustomColor(0xCCCCCC)
for (cc in a..b) {
   // do things
}

3. A Quick Look into IntRange

Simply put, we’ll need to implement Comparable, Iterable, and ClosedRange. From our previous article, we already know we’ll have to implement Comparable.

For the other two interfaces, let’s dive into the IntRange class declaration for some hints:

public class IntRange(start: Int, endInclusive: Int) : 
  IntProgression(start, endInclusive, 1), ClosedRange<Int> 

And then, IntProgression‘s declaration shows that it implements Iterable<Int>:

public open class IntProgression : Iterable<Int>

So, we’re going to want to do something similar to make this work.

4. ColorRange Class

Like IntRange, let’s create a ColorRange class.

For our purposes, we’ll skip mimicking IntProgression, too, since we’re okay with having a default step of 1. This will simplify things a bit and allow us to simply implement both ClosedRange and Iterable directly:

class ColorRange(override val start: CustomColor,
  override val endInclusive: CustomColor) : ClosedRange<CustomColor>, Iterable<CustomColor>{

    override fun iterator(): Iterator<CustomColor> {
        return ColorIterator(start, endInclusive)
    }
}

For our implementation of iterator(), we’ll return a ColorIterator class that will do the heavy lifting of actually stepping through the range.

Because ColorRange implements the ClosedRange<T: Comparable<T>> interface, we have to implement the compareTo method on CustomColor class:

override fun compareTo(other: CustomColor): Int {
    return this.rgb.compareTo(other.rgb)
}

5. ColorIterator Class

ColorIterator is the last piece of the puzzle:

class ColorIterator(val start: CustomColor, val endInclusive: CustomColor) : Iterator<CustomColor> {

    var initValue = start

    override fun hasNext(): Boolean {
        return initValue <= endInclusive
    }

    override fun next(): CustomColor {
        return initValue++
    }
}

Note that initValue is of type CustomColor. So, to mutate it with the ++ operator, we’ll need to add the inc() method to CustomColor as well:

operator fun inc(): CustomColor {
    return CustomColor(rgb + 1)
}

6. Using the Custom Range

We’re almost there!

Since we’re defining our custom range, CustomColor class must implement the rangeTo method. The rangeTo method will allow us to iterate over our range using the .. operator, kind of like how adding inc allows us to use the ++ operator.

Let’s check out the final product:

class CustomColor(val rgb: Int): Comparable<CustomColor> {

    override fun compareTo(other: CustomColor): Int {
        return this.rgb.compareTo(other.rgb)
    }

    operator fun rangeTo(that: CustomColor) = ColorRange(this,that)

    operator fun inc(): CustomColor {
        return CustomColor(rgb + 1)
    }
}

And that’s all we need!

Finally, let’s see how this all works together, by using a range of our CustomColor class:

@Test
fun assertHas10Colors(){
    assertTrue {
        val a = CustomColor(1)
        val b = CustomColor(10)
        val range = a..b
        for (cc in range) {
            println(cc)
        }
        range.toList().size == 10
    }
}

In this test, we’ve defined a range variable and used to iterate through the CustomColor objects, as well as transform it into a list.

Let’s see another example of using the standard contains method on the range:

@Test
fun assertContains0xCCCCCC(){
    assertTrue {
        val a = CustomColor(0xBBBBBB)
        val b = CustomColor(0xDDDDDD)
        val range = a..b
        range.contains(CustomColor(0xCCCCCC))
    }
}

7. Conclusion

Kotlin has a native implementation of range for Int, Long and Char values. In this article, we learned how to implement a range on a custom class.

As always, the code is available on GitHub.