The Ultimate Breakdown of Kotlin Coroutines. Part 5. Suspend Functions.

Previous parts covered suspend lambdas and the machinery, which makes coroutines tick - suspend and resume. Also, I covered continuations, state-machines and interception. There is only one part remaining - suspend functions. I will omit coroutines inlining, which despite having complicated implementation, has an easy mental model - inlining does not affect behavior of coroutines. All the inlining is done prior to state-machine transformation. Finally, I will omit optimizations for the same reason - they do not change the behavior.

Suspend Functions

As explained in the continuation-passing style section, every suspending function's signature is changed: the compiler adds a continuation parameter and changes the return type to Any?.

The tricky part becomes when we try to make it suspendable, in other words, when we build a state machine and generate a continuation. Unlike suspend lambdas, we cannot reuse an existing class for the continuation, since the suspend function can be static, or there can be several functions inside one class.

One way we can solve the issue is to turn the suspend function into a suspend lambda somewhat. We could generate a suspend lambda with code of the suspend function and inside the function call the lambda. For example, when we have a function like

suspend fun test() {
    suspendMe()
    suspendMe()
}

we could generate code like

val test$1: suspend () -> Unit = {
    suspendMe()
    suspendMe()
}

suspend fun test() {
    test$1()
}

As one can see, these two pieces of code are semantically identical. That, by the way, is how JS and Native back-ends generate suspending functions. Furthermore, in JVM, we also used to do this, but not anymore.

The reason why in JVM we do not do this anymore is stack-traces. If we did copy the body of the suspend function to the lambda, the stack-trace would look like

suspendFun1$1.invokeSuspend
suspendFun2$1.invokeSuspend
suspendFun3$1.invokeSuspend
suspendFun4$1.invokeSuspend
suspendFun5$1.invokeSuspend

but we want it to look like

suspendFun1
suspendFun2
suspendFun3
suspendFun4
suspendFun5

thus, instead of moving the function body to lambda, we keep it in the function and build the state-machine there. However, we also keep the 'lambda', so we store all spilled variables there and the label. This 'lambda' is called continuation, and it is, essentially, the state of the coroutine. Unlike suspend lambdas, we split the state (and call it continuation) and the state machine for suspending functions.

Start

Nevertheless, there is another problem. To properly support the completion chain, we need to create the continuation and store the continuation parameter in the completion field. We also need to support resuming the coroutine, i.e., we need to get label and spilled variables from the continuation. So, we need to distinguish these two cases: starting anew and continuing previously suspended execution.

The easiest way to do this is to check for the type of continuation parameter. So, the function preamble will look like:

fun test($completion: Continuation<Unit>): Any? {
    val $continuation =
        if ($completion is test$1) $completion
        else test$1($completion)
    // state machine
}

As long as we generate distinct continuation types for each suspending function, the trick with the check allows us to distinguish these two cases.

However, we have a third case: recursion. When we recursively call the function, the type of the continuation parameter is the same, as if we just resumed. So, there three possible calls of the function:

1. direct call from another suspend function or suspend lambda
2. resumption
3. recursion

So, we need to store at least one another bit of information. We use sign bit of label field for this. Thus, the prefix of the function looks like

fun test($completion: Continuation<Unit>): Any? {
    val $continuation =
        if ($completion is test$1 && $completion.label < 0) {
            $completion.label.sign_bit = 0
            $completion
        } else test$1($completion)
    // state machine
}

here, we assume that the sign bit is unset in recursive calls, while the continuation class sets it during the resume process. So, let us see how we resume and set the bit.

Resume

As we dealt with starting a suspend function and creating a coroutine (in a broad sense), we can tackle the resume process. As explained earlier, when a coroutine (in a narrow sense) suspends, it returns COROUTINE_SUSPENDED. Thus, among the three essential processes of coroutines: creation, suspension, and resumption, there is only the latter left.

In BaseContinuationImpl.resumeWith we call invokeSuspend. So, inside of invokeSuspend, we call the function and pass this as the continuation parameter:

fun invokeSuspend(result: Result<Any?>): Any? {
    test(this)
}

However, we need to set sign bit of the label as well:

fun invokeSuspend(result: Result<Any?>): Any? {
    this.label.sign_bit = 1
    test(this)
}

Let us change the function to call another function that returns a value:

val c: Continuation<Int>? = null

fun suspendInt(): Int = suspendCoroutine {
    c = it
}

suspend fun test() {
    val i = suspendInt()
    println(i)
}

fun main() {
    builder {
        test()
    }
    c?.resume(42)
}

When we run the example, we get 42 printed. Meaning, that the result is somehow passed to the function. The only place we can pass it is invokeSuspend. Also, there is only the continuation parameter of the function. Thus, we need to put the result to the continuation object itself:

fun invokeSuspend(result: Result<Any?>): Any? {
    this.result = result
    this.label.sign_bit = 1
    test(this)
}

then, get the result from the continuation object in the function:

fun test($completion: Continuation<Unit>): Any? {
    val $continuation =
        if ($completion is test$1 && $completion.label < 0) {
            $completion.label.sign_bit = 0
            $completion
        } else test$1($completion)
    val $result = $continuation.result
    // state machine
}

Variables spilling is the same regardless, whether it is a lambda or a function. However, we spill the variables to the continuation object.

JVM: Parameters

Let us now have a look into how we deal with suspend function parameters. We do not generate fields for them, since a lambda uses them just to pass the arguments from invoke to invokeSuspend. We do not need them for suspending functions: the arguments are locals; thus, we reuse local variables spilling for them.

Nevertheless, keeping the parameters in the function signature breaks resumption. There is not enough information in continuation's invokeSuspend to pass them as they were before or as they are now. So, we just put nulls for reference types and zeroes for primitives. That means we cannot generate nullability checks at the beginning of the function, so they must be generated at the beginning of the first state, to which we cannot resume.

For example, if we change the test function to accept an argument:

suspend fun dummy() {}

suspend fun test(a: Int) {
    dummy()
    dummy()
}

its continuation's invokeSuspend becomes something like

fun invokeSuspend(result: Result<Any?>): Any? {
    this.result = result
    this.label.sign_bit = 1
    test(0, this)
}

Local Suspend Functions

Local functions are weird: the way the compiler generates them is different in back-ends. Local suspend functions are even stranger.

Old JVM: Unerased Suspend Lambdas

Old JVM back-end generates local functions as lambdas. Thus, suspend local functions are generated as suspend lambda:

fun main() {
    suspend fun local(i: Int) {}
}

is generated as something like:

fun main() {
    val local: suspend (Int) -> Unit = {}
}

Doing so allows us to reuse the logic of captured variables and simplify the logic of code generation. However, because of the limitations of old BE, its create and invoke are unerased. In other words, the compiler duplicates them, generating unerased and erased copies. Unerased copy accepts typed parameters, and it contains the logic of lambda's invoke or create. The other accepts only Any? parameters since they override supertype's functions, and delegates call to an unerased copy.

JVM_IR: Static Functions

On the other hand, JVM_IR generates local functions as static functions with captured variables put as first parameters. Thus, suspend local functions are generated as static functions as well. That reduces code size and method count and enables tail-call optimization.

An example:

fun main() {
    val aa: Long = 1
    suspend fun local(i: Int) {
        println(aa)
    }
}

is generated as something like

suspend fun main$1(aa: Long, i: Int) {
    println(aa)
}

fun main() {
    val aa: Long = 1
}

Summary

This concludes the breakdown of Kotlin coroutines. I hope it helped you get the grasp of complexity of simple-on-the-surface feature. And I hope you found answers to the questions you might have had. If not, you are welcome in #coroutines channel of official slack: kotlinglang.slack.com.