Fun with Kotlin: Creating a DSL for Constructing Binary Trees

Photo by Gilly Tanabose on Unsplash

In this post, we’re going to implement a DSL for constructing binary trees declaratively. We will not jump directly to the DSL implementation but rather start with the most meticulous (and unreadable) approach and gradually move towards our DSL. Let’s start!

For the rest of this article, we use the following binary tree as an example.

First Approach

The first thing we need is to define the class for a binary tree node. We can do it like this:

class BinaryTreeNode<T>(
    var value: T,
    var left: BinaryTreeNode<T>? = null,
    var right: BinaryTreeNode<T>? = null
)

Now we can write the following code to create a binary tree.

fun main() {
    val root = BinaryTreeNode("A")
    val b = BinaryTreeNode("B")
    val c = BinaryTreeNode("C")
    val d = BinaryTreeNode("D")
    val e = BinaryTreeNode("E")
    val f = BinaryTreeNode("F")

    root.left = b
    root.right = c

    b.left = d
    b.right = e

    c.right = f
}

As can be seen, not readable at all! When someone looks at the above code they should carefully follow the code to be able to have an idea of the structure of the tree.

Second Approach

For our second approach, we use apply scope function to make our tree construction code more readable. We can write a code like the following.

fun main() {
    val root = BinaryTreeNode("A").apply {
        left = BinaryTreeNode("B").apply {
            left = BinaryTreeNode("D")
            right = BinaryTreeNode("E")
        }
        right = BinaryTreeNode("C").apply {
            right = BinaryTreeNode("F")
        }
    }
}

It’s obvious that this approach is far more readable than our previous approach. For instance, we can easily see that nodes “D” and “E” are the left and right children of node “B” respectively.

To make the above code a bit more readable we can define a small function node to create a new instance of BinaryTreeNode. Using the node function, we will end up with the following code.

fun main() {
    val root = node("A").apply {
        left = node("B").apply {
            left = node("D")
            right = node("E")
        }
        right = node("C").apply {
            right = node("F")
        }
    }
}

fun <T> node(value: T) = BinaryTreeNode(value)

The change we made was a small one but it helped to have a slightly more readable code.

Using a DSL

So far we have seen two approaches to create a binary tree. Now it’s time to create our own DSL for creating such a tree. What we want to achieve is like the following code.

fun main() {
    val root = binaryTree("A") {
        left("B") {
            left("D")
            right("E")
        }
        right("C") {
            right("F")
        }
    }
}

In the next section, we will see how we can implement such a DSL for creating binary trees.

Implementing the DSL

We will start with defining a function named binaryTree which we can call to start creating a tree.

fun <T> binaryTree(rootValue: T, content: BinaryTreeBuilder<T>.() -> Unit = {}): BinaryTreeNode<T> {
    val root = BinaryTreeNode(rootValue)
    val binaryTreeBuilder = BinaryTreeBuilder(root)
    binaryTreeBuilder.content()
    return root
}

Let’s see what the above function does. The first parameter of binaryTree is the value of the tree root. The value of a node could be of any type (Int, String, custom classes, you name it). The second parameter is a lambda with receiver (we will talk about the BinaryTreeBuilder later on). The return value of the function is the root of the binary tree.

In the body of the binaryTree function we do several things:

• Create the root node of the binary tree (root).
• Create a new instance of BinaryTreeBuilder and pass root to its constructor.
• Call the content lambda on the BinaryTreeBuilder instance
• Return the root node as the root of the binary tree

Using Kotlin’s trailing lambda syntax, we can call the above function like below:

binaryTree("A") {
}

Inside the lambda is where we can start adding nodes to the tree. The lambda has a receiver with typeBinaryTreeBuilder that allows adding nodes to the tree.

Let’s see the implementation of BinaryTreeBuilder.

class BinaryTreeBuilder<T>(root: BinaryTreeNode<T>) {
    private var currentParent = root

    fun left(value: T, content: () -> Unit = {}) {
        val parent = currentParent
        val newNode = BinaryTreeNode(value)
        parent.left = newNode
        currentParent = newNode
        content()
        currentParent = parent
    }

    fun right(value: T, content: () -> Unit = {}) {
        val parent = currentParent
        val newNode = BinaryTreeNode(value)
        parent.right = newNode
        currentParent = newNode
        content()
        currentParent = parent
    }
}

Let’s see how the above code works.

Looking at the above code we can see that, at any given time, currentParent keeps track of the current node that we want to add a node to it as its child.

There are also two methods left and right that is used to add a left and a right node as the children of the current node respectively. Let’s see how function left works:

• A new node (newNode) is created and set as the left child of the current parent node (currentParent),
• Temporarily newNode is set as the new current parent and the content lambda executes. This effectively causes whatever nodes added in the lambda to be the children of newNode.
• After executing the content lambda currentParent is restored to its previous value.

The right function works in a similar way except that newNode is added as the right child of currentParent.

That’s it. Our DSL is ready to use!

Where I use this DSL

Every once in a while I spend some time solving LeetCode problems. When the problem is a binary tree one, after implementing the algorithm I need to create some binary trees as test cases to test my algorithm against them. I use the DSL we just developed to easily create those test cases!

You can access the source code of this post from here.