Kotlin Extension Functions

Back in the early 1990s, I was a big fan of a US science-fiction TV show called Quantum Leap. If you’re not familiar with this masterpiece of casual popular sci-fi, a summary of the show’s premise was provided at the start of each episode:

Theorizing that one could time travel within his own lifetime, Dr. Sam Beckett stepped into the Quantum Leap accelerator, and vanished. He woke to find himself trapped in the past, facing mirror images that were not his own, and driven by an unknown force to change history for the better.

What followed was a series of episodic delights in which Dr Beckett found himself transplanted into the body of some character in history, whose life was impacted by a bad decision of their own or somebody close to them. Beckett’s job was to fix the mistake, change the course of history and move on to the next story.

While such adventure generating time travel technology has so far not arrived. The Kotlin programming language gives us something close. The extension function. Using this feature of the Kotlin language we can write something that looks like this:

fun String.newFunction() = //something here

When this code is passed through the Kotlin compiler, an avatar is created at the offices of Sun Microsystems somewhere around 1995, where it will whisper in the ear of the developer working on the String class that would form part of the Java standard library and convince them to add this function to the set of those provided in the class definition.

This magic of time travel intervention now complete, we are free to use this new temporally injected function in our own code like this:

"Hello World".newFunction()

This, at least, is how we are encouraged to understand the feature. Even the name, “extension function” suggests this interpretation. In general, it is a fairly good metaphor for many typical uses of the extension function system. Most of the time we are simply augmenting an existing type with additional, helpful functionality.

Where it can break down as a metaphor is in cases when extension functions are used as first-class entities. Consider:

fun myFunkyFun(argument: String.(Int) -> Unit): List.(Double) -> Character

This, albeit convoluted example demonstrates the case of a function that both accepts and returns an extension function type. The metaphor of these functions being “monkey patched” into existing types is less useful for conceptualising this example; and understanding how and when to use these types of constructions is critical for applying some of the more advanced patterns that extension functions allow us to create, such as type-safe-builders.

The purpose of this post is to encourage another way to understand extension functions. One that, I believe, supports a better understanding of how to use them as powerful building blocks for more fluent and maintainable code.

Peeking behind the curtain

The standard Java toolkit comes with the useful, but slightly scary javap tool. The javap tool allows us to disassemble our compiled class files back into the something resembling the code that produced them. It can be particularly useful for inspecting the behaviour of features of Kotlin that do not exist in Java as you can create the class file with the Kotlin compiler and then use javap to inspect what the equivalent Java source code might look like.

Let’s take a Kotlin example containing an extension function and a regular function and compare the results when disassembled with javap.

fun simpleFunction(string: String) = "Hello, $string"
fun String.extensionFunction() = "Hello, $this"

When we compile this code and inspect it with javap we see:

public static final java.lang.String simpleFunction(java.lang.String);
public static final java.lang.String extensionFunction(java.lang.String);

Exactly the same 🤯.

This is the key to understanding extension functions. If we abandon the purely object-centric view of the world that languages like Java questionably invite us to hold, and instead look at them with a more functional brain, we begin to see that the receiver of these functions is just another argument.

Kotlin already knows this stuff

A good way to solidify this concept is to demonstrate that the Kotlin compiler considers the the following types as entirely interchangeable:

String.() -> String
(String) -> String

Let’s play with this idea:

fun simpleFunction(string: String) = "Hello, $string"
fun String.extensionFunction() = "Hello, $this"

fun expectsSimple(function: (String) -> String) =
    function("simple")

fun expectsExtension(function: String.() -> String) =
    "extension".function()

fun main() {
    println(expectsSimple(String::extensionFunction))
    println(expectsExtension(::simpleFunction))
}

Here, we define two additional functions. The first, expectsSimple, expects to be provided with a simple function that takes a string and returns a string, calling it with the argument “simple” and returning the result. The second, expectsExtension, expects to be provided with a String extension function and calls that extension function on the value of “extension”, returning the result.

In the main method, we flip this script. We pass the extension function to expectsSimple and the simple function to expectsExtension. On running this code we get:

Hello, simple
Hello, extension

The Kotlin compiler treats the two different types as being equal. It knows that you will need to pass a string value to both function types. The difference is whether you pass that value in an argument position, or in a receiver position, but the input data types and the output data types are exactly the same in both cases.

The functional viewpoint

Armed with this information, there is another way we can start to look at extension functions. Indeed, functions in general. In languages like Kotlin, Java, Javascript and many others, every function that you call has a “context argument”. The context argument is the value that can be referred to using this within the function body. And, as we know, this can be omitted when accessing it’s members. this.someFunction() can be written as someFunction(). Extension functions are just functions where the type of the context argument is defined explicitly by the programmer.

Building the builder

We can now put all this together and create a simple type-safe-builder.

Let’s assume we want to build a document. Documents for our purposes are defined as having a string title and a list of paragraphs that each have their own string title and some string content. In other worse, they can be defined by these data classes…

data class Paragraph(val title: String, val content: String)
data class Document(val title: String, val paragraphs: List<Paragraph>)

While, in a simple case like this, constructing these data classes directly would hardly be a chore, we have further decided that we’d like to support a builder syntax that works as follows…

val doc = document("My title") {
    paragraph("Paragraph 1") {
        "This is my first paragraph"
    }
    paragraph("Paragraph 2") {
        "This is my second paragraph"
    }
}

How do we go about implementing the infrastructure to enable this DSL-like syntax to build our document types?

Looking at the code above, we can clearly see there is a need for a document function and that this function expects two arguments. The first argument is a simple string, the document title. The second argument is a function, passed here as a lambda. However, the lambda has a very particular requirement. There must be a context object that provides for another function named paragraph that also takes two arguments. Again, a simple string for the title, and additionally a function that returns a string to be used for the content.

Let’s define this context type as a class:

class ParagraphBuilder {

    val paragraphs = mutableListOf<Paragraph>()
    
    fun paragraph(title: String, contentBuilder: () -> String) {
        paragraphs.add(Paragraph(title, contentBuilder()))
    }
}

Finally, we can define our document function by expressing that it’s second argument - the function - is an “extension function” of our ParagraphBuilder class. What we really mean by “extension function” is, “a function that uses a ParagraphBuilder as the type of it’s context object”. It will look like this:

fun document(title: String, buildParagraphs: ParagraphBuilder.() -> Unit): Document {
    val paragraphBuilder = ParagraphBuilder()
    paragraphBuilder.buildParagraphs()
    return Document(title, paragraphBuilder.paragraphs)
}

The document function creates a ParagraphBuilder to serve as the builder for this particular execution of the function. It then invokes our passed function argument “on” this builder, which is to say, invokes it by providing it with the actual instance of the context object that the function should use while executing and finally constructs the new document using the passed title and the paragraphs that were built during the execution of the buildParagraphs function.

While this is a very simple case, it is easy to now see how such a pattern can be extended into more complicated builders of more advanced structural data. All the while, maintaining the type rules about which values can be added in which contexts.

At last, the conclusion

The purpose here is to invite the conception of the Kotlin extension function as something more than “magically adding functions to existing types”. In reality, what we are doing is defining functions that run with a specific context that is different from the default. By conceiving of extension functions as functions that run with an explicitly specified context, more advanced uses of these types of functions are unlocked.

Go forth and extend.