I am looking to see if there is any design document reading with respect to how Kotlin’s coroutines are implemented as state machines under the hood specifically with respect to looping.
But so far, I don’t see much detail about how the compiler state machines are created when looping is present in the code. Either infinite loops or bounded loops. I’m hoping there is reading material that already exists or that someone can explain a few examples of how the CPS state machine is created when it has to consider a loop of some sort. If I have to deep dive into code perhaps someone can point in the direction of where this code lives.
I’m also willing to contribute such documentation back to the repo to pay any explanation forward.
Unfortunately, I don’t have any documentation to read, but what is you concern regarding loops? In simple words, when Kotlin compiles a suspend function, then for each suspension point (call to another suspend function), it adds code for storing local variables and jumping out of the function and another code fragment to jump into this specific point and load variables back. I don’t think loops are any different here, but maybe I miss something.
I don’t have any concerns other than just curiosity of how this works. What you said makes sense to me in the general sense for how the compiler rewrites suspendable functions (coroutines).
I guess what I’m getting at is how a function that has a loop with some kind of “yielding” statement is handled during the transformation phase?
Example but this is just pseudocode (not Kotlin):
func foo(){
for i=0; i<3; i++ {
println("do some useful work..")
yield()
}
}
The part that I would like to understand is how the function above is transformed when there is a loop present. Not only does the function have to suspend due to the yield, but it has to do this 3 times exactly to satisfy the loop.
Just to be clear, I’m just curious there’s no problem per se.
switch (loadLocation()) {
case 1 -> {
loadVars(i)
goto :suspension1
}
i = 0
:loop
println("do some useful work..")
saveLocation(1)
saveVars(i)
yield()
if (isSuspended()) {
return
}
:suspension1
i++
if i < 3 :loop
return
Some notes:
I don’t go into details of saving and loading, checking if suspended, etc., but it is generally utilizing continuation objects. The real bytecode is a little more complicated.
As far as I remember it really saves the state before yield() and unconditionally. To be honest, I’m not sure why it doesn’t store it inside if (suspend) {, because if we don’t suspend then saving is a wasted overhead. Maybe this is to correctly handle exceptions, but I don’t know.
val fibonacciSeq = sequence {
var a = 0
var b = 1
yield(1)
while (true) {
yield(a + b)
val tmp = a + b
a = b
b = tmp
}
}
Now, next to the version dropdown select JS. A link will show up (called: Generate JS) that allows you to see the generated JavaScript state-machine. For whatever reason I can’t get it to show JVM-based output but the JS state-machine is still very helpful and demonstrates the pattern utilized.
Here is some example output for the above code (showing specifically the doResume function:
Coroutine$fibonacciSeq$lambda.prototype.doResume = function () {
do
try {
switch (this.state_0) {
case 0:
this.local$a = 0;
this.local$b = 1;
this.state_0 = 2;
this.result_0 = this.local$$receiver.yield_11rb$(1, this);
if (this.result_0 === COROUTINE_SUSPENDED)
return COROUTINE_SUSPENDED;
continue;
case 1:
throw this.exception_0;
case 2:
this.state_0 = 3;
continue;
case 3:
this.state_0 = 4;
this.result_0 = this.local$$receiver.yield_11rb$(this.local$a + this.local$b | 0, this);
if (this.result_0 === COROUTINE_SUSPENDED)
return COROUTINE_SUSPENDED;
continue;
case 4:
var tmp = this.local$a + this.local$b | 0;
this.local$a = this.local$b;
this.local$b = tmp;
this.state_0 = 3;
continue;
default:
this.state_0 = 1;
throw new Error('State Machine Unreachable execution');
}
} catch (e) {
if (this.state_0 === 1) {
this.exceptionState_0 = this.state_0;
throw e;
} else {
this.state_0 = this.exceptionState_0;
this.exception_0 = e;
}
}
while (true);
};
A few observations of my own
Every time a suspendable function is called such as yield, the state-machine will have more case statements to handle this.
Whenever variables live across suspension boundaries they must be hoisted and tracked in the continuation itself and not in the function. This is why you see references to this.local$a or this.local$b.
Case statement number 0 seems to be related to handling initializing of the function
Case statement number 1 always seems to be reserved for bubbling up exceptions.
The default case should never be reachable in normal operation.
Please be aware it is different in JVM. I sometimes looked into the generated bytecode and it is not this typical state machine as above, where we handle states in a loop and each can choose where to switch to. It is more similar to the regular bytecode for the same source code, but with added saving, loading, jumping out and in. We don’t process states in a loop, but if yield(1) doesn’t suspend, the execution naturally flows into while (true) {. I believe we run through the state-switch only when resuming.
But generally yes, you are correct. Every new call to a suspend function creates a new suspension point, which creates a new state. Yes, all variables that could be needed later, have to be saved into the continuation before suspending. And yes, from my experience state 0 is the beginning of a function. But when I looked into the bytecode of the above function for JVM, state 1 was just a regular state related to yield(1) suspension point.
Also passing of continuations is handled a little differently in JVM. We check if the received continuation is “ours” (we check its type). If it is of unknown type, then we assume this is the initial call to the function and the passed continuation is a continuation of the caller. Then we initialize a new continuation. If the passed continuation is of our type, we assume this is a resumption. I’m not exactly sure how does it handle the case of a function calling itself, but it doesn’t seem hard to solve.
Thanks @broot for the clarification. It would be great to compare the JVM output to the JS output in the playground to see how they differ but I completely understand the implementations may be different to a varying degree.
I definitely appreciate all the additional context you’ve provided as well on this topic!
I can provide you the bytecode for the above, but it is pretty long and much harder to read than JavaScript. This is what I got while using Kotlin 1.8.20:
When targeting the JVM, all loops are translated into (conditional or unconditional) jumps anyway. Therefore resuming in the middle of a loop is no different to resuming at any other place in the code. Before every suspension point, the compiler has added code that stores away the transient state, which includes “loop variables”, together with a label for the current position within the method. If the program needs to resume it will jump to the stored label, where more code has been added that restores the remaining state. The program can the continue where it left off.