Playing around with Kotlin Sealed Classes
Tags: kotlin, jackson, java
Published: 2020-09-24 06:50:15

While working on my Master's thesis, I recently had to configure probabilities. These values were organized in a tree-like structure. Because of this, I thought of using inheritance - and with Kotlin already in place, using sealed classes seemed at least like an interesting opportunity to me.

> **Disclaimer:** This post is just for fun and for showing some cool Kotlin-specific things. In the end, using Jackson to deserialize input into data classes is probably more effective.

## Sample

Let's say, we want to configure the likeliness in percent of some actions using YAML for instance. This could look like the following:
```yaml
Configuration:
    Greeting:
        SayHello: 70
        WaveAt: 80
        Hug: 5
    Talking:
        DoSmallTalk: 90
        Insult: 1
```
As you can see, there are two subsections. However, this can easily get an arbitrary depth.

Before we'll reproduce that structure, we should clarify what sealed classes are.

## Sealed Classes

The [official documentation](https://kotlinlang.org/docs/reference/sealed-classes.html) calls them extensions to enum classes, used for representing restricted class hierarchies.

Sealed classes are declared using the keyword `sealed`, are abstract and must have only private constructors.

Simply put, a sealed class ensures, that **all** possible subtypes are known at compile time. This comes especially handy, when used in conjunction with Kotlin's `when` clause:
```kotlin
sealed class TruthOrDare
class Truth(val question: String) : TruthOrDare()
class Dare(val task: String) : TruthOrDare()

fun nextTurn(input: TruthOrDare) = when(input) {
    is Truth -> println("Answer the following question: ${input.question}")
    is Dare -> println("Your task is: ${input.task}")
    // the `else` clause is not required because we've covered all the cases
}
```

## Sealed Classes in Action

Let's come back to the original task: using sealed classes for parsing configuration.

The structure of the above YAML snippet can be represented as follows:
```kotlin
sealed class Configuration {
    
    sealed class Greeting : Configuration() {
        
        object SayHello : Greeting()
        object WaveAt : Greeting()
        object Hug : Greeting()
        
    }
    
    sealed class Talking : Configuration() {
        
        object DoSmallTalk : Talking()
        object Insult : Talking()
        
    }

}
```

We have nested sealed subclasses and `object`s (a shortcut for [Singletons](https://en.wikipedia.org/wiki/Singleton_pattern) in Kotlin). We'll see later, why we are using objects here.

However, yet there is nothing configured. So how do we use this structure?

The answer is: [Reflection](https://kotlinlang.org/docs/reference/reflection.html).

### Digression
Before we come to the actual structure though, I would like to introduce a small wrapper class, to encapsulate our findings:
```kotlin
class Probabilities private constructor(
        private val backingMap: MutableMap<Configuration, Int> = mutableMapOf()
) : Map<Configuration, Int> by backingMap {

    constructor(configuration: Probabilities.() -> Unit) : this() {
        configuration()
    }

    infix fun Configuration.withProbabilityOf(percent: Int) = backingMap.put(this, percent)

    override fun toString(): String = backingMap.entries.joinToString { "${it.key::class.simpleName} = ${it.value}" }
}
```
In this class you can see some other cool Kotlin features: [**implementation by delegation**](https://kotlinlang.org/docs/reference/delegation.html), [**receiver functions**](https://kotlinlang.org/docs/reference/lambdas.html#function-types) and [**infix functions**](https://kotlinlang.org/docs/reference/functions.html#infix-notation).

Implementing the `Map` interface can be a pain in Java. Kotlin provides a smart solution with **implementation by delegation**. All you need to do, is to use the `by` keyword in conjunction with a value, that already implements the interface. So in this case we simply delegate all `Map`-specific operations to a `MutableMap`, which is created by the default constructor.

To make configuration of the `Probabilities` class easy, we provide a second constructor taking a **receiver function**. This means, that the `this` keyword of the provided lambda points to an instance of `Probabilities` class.

Another **receiver function** is also provided by `withProbabilityOf`. However, this is also an **infix function**, which is marked by the `infix` keyword and enables you to write something like:
```kotlin
SayHello withProbabilityOf 99
```

## Parsing
Finally, here is the code to parse e.g. a `Map` of `String` and `Any`:
```kotlin
fun fromMap(configuration: Map<String, Any>) = Probabilities { // (1)
    @Suppress("UNCHECKED_CAST")
    fun parseFor( // (2)
            configuration: Map<String, Any>, // (3)
            parent: KClass<out Configuration>?,
            clazz: KClass<out Configuration>
    ) {
        // (4)
        if (configuration.containsKey(clazz.simpleName) && parent?.isSuperclassOf(clazz) != false) {
            when (val value = configuration[clazz.simpleName]) {
                is Map<*, *> -> { // (5)
                    clazz.sealedSubclasses.forEach { subclass ->
                        parseFor(value as Map<String, Any>, clazz, subclass)
                    }
                }
                is Int -> { // (6)
                    if (clazz.objectInstance == null)
                        throw RuntimeException("${clazz.simpleName} should be an object")

                    clazz.objectInstance!! withProbabilityOf value
                }
                else -> throw RuntimeException("unknown property ${clazz.simpleName}")
            }
        }
    }

    parseFor(configuration, null, Configuration::class)
}
```

Some explanation (mind the comments in code):

 * **(1)** we construct a new instance of `Probabilities` by making use of one of Kotlin's greatest features: if the last parameter of any function or constructor is a lambda function, we can omit parentheses and only use curly braces
 * **(2)** another beautiful thing are `local functions`: since we don't need the recursive function `parseFor` anywhere else, we simply declare it inside our `fromMap` function
 * **(3)** the parameters of our `parseFor` function are:
   * the subtree of the configuration
   * the assumed parent as specified by the configuration
   * the current class as specified by the configuration's property
 * **(4)** if somewhere on top level of the current subtree the expected class is found and the assumed parent matches the real superclass, then we use `when` on the property value's type
 * **(5)** if it's a `Map`, apply `parseFor` on each sealed subclass
 * **(6)** if it's an `Int` and `clazz` is an `object`, we configure the probability to be that value

In the above code, you might see, why having leaves being `object`s is a good idea: there will always be exactly one instance. So they are perfect for being used as a key in a map.

## Sweet. - But why?
To be honest, all of the above might seem like an over-engineered construct. Compared to the [example using sealed classes](#sealed-classes-in-action), an equivalent structure with data classes would be similar and could be easily used with e.g. Jackson out of the box:
```kotlin
data class Configuration(
    val greeting: Greeting?,
    val talking: Talking?
) {

    data class Greeting(
        val sayHello: Int?,
        val waveAt: Int?,
        val hug: Int?
    )

    data class Talking(
        val doSmallTalk: Int?,
        val insult: Int?
    )
}
```
This is the great thing about Kotlin: you have so many language constructs that help you build your software the *best* way, whatever this means. Therefore, it's necessary to play around with and get to know these possibilities.

## Using Jackson

To make use of the concept of using sealed classes in conjunction with Jackson, we need to provide a custom deserializer. Since we already have the mapping logic however, this is an easy task:
```kotlin
class ProbabilitiesDeserializer : StdDeserializer<Probabilities>(Probabilities::class.java) {

    override fun deserialize(p: JsonParser?, ctxt: DeserializationContext?) =
            p?.readValueAs<Map<String, Any>>(object : TypeReference<Map<String, Any>>() {})
                    ?.let { fromMap(it) }

}
```

For this and the ability to parse YAML we need the following dependencies:
 * `com.fasterxml.jackson.module:jackson-module-kotlin:2.11.2`
 * `com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.11.2`

Finally, we are ready to create an `ObjectMapper` and read our YAML configuration:
```kotlin
val mapper = ObjectMapper(YAMLFactory()).apply {
    registerModule(SimpleModule().addDeserializer(Probabilities::class.java, ProbabilitiesDeserializer()))
    registerModule(KotlinModule())
}
val probabilities = mapper.readValue<Probabilities>(yamlInput)
println(probabilities.toString())
println("Say hello with probability of ${probabilities[SayHello]} percent.")
```

The result will be: 
```
SayHello = 70, WaveAt = 80, Hug = 5, DoSmallTalk = 90, Insult = 1
Say hello with probability of 70 percent.
```

That's it.

## Bonus: Ensuring all properties are set

One last thing I wanted to add, is how you enforce all `object`s to be configured.

All we have to do is to find all the leaves of our configuration tree and compare them to the input's content:
```kotlin
val leaves = leavesOf(Configuration::class)
fun leavesOf(baseClass: KClass<out Configuration>): List<KClass<out Configuration>> =
    if (!baseClass.isSealed) {
        listOf(baseClass)
    } else {
        baseClass.sealedSubclasses.flatMap(::leavesOf)
    }

fun Probabilities.ensureAllActionsCovered() {
    val keys = keys.map { it::class }
    val unconfigured = leaves.filter { leaf -> !keys.contains(leaf) }
    if (unconfigured.isNotEmpty())
        throw RuntimeException("Unconfigured leaves: ${unconfigured.joinToString()}")
}
```

Finally, we've seen a last cool Kotlin feature: [**extension functions**](https://kotlinlang.org/docs/reference/extensions.html#extension-functions). They allow us to add some functionality to an otherwise closed class. Like a bonus in a way.

You call them as if they were a class method:
```
probabilities.ensureAllActionsCovered()
```

## Closing notes

Thanks for reading, I hope you liked it. You can find all the code in a [kscript](https://github.com/holgerbrandl/kscript) on [Github as a Gist](https://gist.github.com/alxgrk/e672b031206481bc462d21788a321804).