Kotlin Coroutine Design with respect to looping

Greetings,

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.

I am aware of these links: coroutine code-gen

and this link: implementation details

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.

Thanks in advance!

-Ralph

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.

3 Likes

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.

As far as I know it will be something like:

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.

2 Likes

Look here: KEEP/coroutines.md at master · Kotlin/KEEP · GitHub

Thank you I was aware of this link but it doesn’t go into much depth about how loops are handled.

For anyone who wants to see how loops are handled:

Go to: play.kotlinlang.org

From there you can type some code in such as:

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.
1 Like

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.

1 Like

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:

  public static final java.lang.Object bbb(kotlin.sequences.SequenceScope<? super java.lang.Integer>, kotlin.coroutines.Continuation<? super kotlin.Unit>);
    Code:
       0: aload_1
       1: instanceof    #28                 // class Test1Kt$bbb$1
       4: ifeq          39
       7: aload_1
       8: checkcast     #28                 // class Test1Kt$bbb$1
      11: astore        6
      13: aload         6
      15: getfield      #32                 // Field Test1Kt$bbb$1.label:I
      18: ldc           #33                 // int -2147483648
      20: iand
      21: ifeq          39
      24: aload         6
      26: dup
      27: getfield      #32                 // Field Test1Kt$bbb$1.label:I
      30: ldc           #33                 // int -2147483648
      32: isub
      33: putfield      #32                 // Field Test1Kt$bbb$1.label:I
      36: goto          49
      39: new           #28                 // class Test1Kt$bbb$1
      42: dup
      43: aload_1
      44: invokespecial #34                 // Method Test1Kt$bbb$1."<init>":(Lkotlin/coroutines/Continuation;)V
      47: astore        6
      49: aload         6
      51: getfield      #38                 // Field Test1Kt$bbb$1.result:Ljava/lang/Object;
      54: astore        5
      56: invokestatic  #44                 // Method kotlin/coroutines/intrinsics/IntrinsicsKt.getCOROUTINE_SUSPENDED:()Ljava/lang/Object;
      59: astore        7
      61: aload         6
      63: getfield      #32                 // Field Test1Kt$bbb$1.label:I
      66: tableswitch   { // 0 to 2
                     0: 92
                     1: 144
                     2: 218
               default: 260
          }
      92: aload         5
      94: invokestatic  #50                 // Method kotlin/ResultKt.throwOnFailure:(Ljava/lang/Object;)V
      97: iconst_0
      98: istore_2
      99: iconst_1
     100: istore_3
     101: aload_0
     102: iconst_1
     103: invokestatic  #56                 // Method kotlin/coroutines/jvm/internal/Boxing.boxInt:(I)Ljava/lang/Integer;
     106: aload         6
     108: aload         6
     110: aload_0
     111: putfield      #59                 // Field Test1Kt$bbb$1.L$0:Ljava/lang/Object;
     114: aload         6
     116: iload_2
     117: putfield      #62                 // Field Test1Kt$bbb$1.I$0:I
     120: aload         6
     122: iload_3
     123: putfield      #65                 // Field Test1Kt$bbb$1.I$1:I
     126: aload         6
     128: iconst_1
     129: putfield      #32                 // Field Test1Kt$bbb$1.label:I
     132: invokevirtual #71                 // Method kotlin/sequences/SequenceScope.yield:(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
     135: dup
     136: aload         7
     138: if_acmpne     172
     141: aload         7
     143: areturn
     144: aload         6
     146: getfield      #65                 // Field Test1Kt$bbb$1.I$1:I
     149: istore_3
     150: aload         6
     152: getfield      #62                 // Field Test1Kt$bbb$1.I$0:I
     155: istore_2
     156: aload         6
     158: getfield      #59                 // Field Test1Kt$bbb$1.L$0:Ljava/lang/Object;
     161: checkcast     #67                 // class kotlin/sequences/SequenceScope
     164: astore_0
     165: aload         5
     167: invokestatic  #50                 // Method kotlin/ResultKt.throwOnFailure:(Ljava/lang/Object;)V
     170: aload         5
     172: pop
     173: aload_0
     174: iload_2
     175: iload_3
     176: iadd
     177: invokestatic  #56                 // Method kotlin/coroutines/jvm/internal/Boxing.boxInt:(I)Ljava/lang/Integer;
     180: aload         6
     182: aload         6
     184: aload_0
     185: putfield      #59                 // Field Test1Kt$bbb$1.L$0:Ljava/lang/Object;
     188: aload         6
     190: iload_2
     191: putfield      #62                 // Field Test1Kt$bbb$1.I$0:I
     194: aload         6
     196: iload_3
     197: putfield      #65                 // Field Test1Kt$bbb$1.I$1:I
     200: aload         6
     202: iconst_2
     203: putfield      #32                 // Field Test1Kt$bbb$1.label:I
     206: invokevirtual #71                 // Method kotlin/sequences/SequenceScope.yield:(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
     209: dup
     210: aload         7
     212: if_acmpne     246
     215: aload         7
     217: areturn
     218: aload         6
     220: getfield      #65                 // Field Test1Kt$bbb$1.I$1:I
     223: istore_3
     224: aload         6
     226: getfield      #62                 // Field Test1Kt$bbb$1.I$0:I
     229: istore_2
     230: aload         6
     232: getfield      #59                 // Field Test1Kt$bbb$1.L$0:Ljava/lang/Object;
     235: checkcast     #67                 // class kotlin/sequences/SequenceScope
     238: astore_0
     239: aload         5
     241: invokestatic  #50                 // Method kotlin/ResultKt.throwOnFailure:(Ljava/lang/Object;)V
     244: aload         5
     246: pop
     247: iload_2
     248: iload_3
     249: iadd
     250: istore        4
     252: iload_3
     253: istore_2
     254: iload         4
     256: istore_3
     257: goto          173
     260: new           #73                 // class java/lang/IllegalStateException
     263: dup
     264: ldc           #75                 // String call to \'resume\' before \'invoke\' with coroutine
     266: invokespecial #78                 // Method java/lang/IllegalStateException."<init>":(Ljava/lang/String;)V
     269: athrow
}
  • 0-59 is initialization, preparing stuff, checking the type of passed continuation, initializing a new one if needed, etc.
  • 66 is a state switch. 0 is the beginning of a function, 1 is yield(1) and 2 is yield(a + b).
  • 92-132 is the beginning of a real code, until the call to yield(1).
  • 135-143 - check if yield(1) suspended and if so, return.
  • 144-170 - loading the state after resuming from state 1 (yield(1))
  • 172-206 - progress further from state 1, until the next suspension point.
  • and so on.
1 Like

You can use IntelliJ Idea to see JVM bytecode and decompile it to Java.

Please consider that every yield suspend the current coroutine, it doesn’t instantiate a new one.

1 Like

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.

2 Likes