Continuation-Passing Style (CPS) — How Kotlin Coroutines Work Under the Hood

Continuation.kt

/**
  This gist demonstrates how Kotlin coroutines internally use continuation-passing style (CPS) and 
  a state machine with switch case, by manually simulating coroutine behavior.
  Coroutines appear linear, but under the hood they are compiled into:
  -> a Continuation object
  -> a label/state machine
  -> a big when {} block that decides where to resume multiple functions that 
     get stitched together during suspension/resumption
**/

/**
 * Tasks need to be executed in the order
 * taskA() -> taskB() -> taskC() -> taskD()
 * and can only be accessed via stateMachinePart1() and stateMachinePart2() functions.
 * Unwrapping what happens when we call coroutineWay()
 * suspend fun coroutineWay()
 * {
 *     taskA()
 *     taskC()
 *     taskB()
 *     taskD()
 * }
 */

const val CONTINUATION_TAG = "ContinuationPassingStyle"

private fun taskA(): String = "A"
private fun taskB(): String = "B"
private fun taskC(): String = "C"
private fun taskD(): String = "D"

interface DemoContinuation {
    fun resumeWith(result: String)
}

fun stateMachinePart1(state: Int, continuation: DemoContinuation?) {
    when (state) {
        1 -> {
            val a = taskA()
            continuation?.resumeWith(a)
        }

        2 -> {
            val b = taskB()
            continuation?.resumeWith(b)
        }
    }
}

fun stateMachinePart2(state: Int, continuation: DemoContinuation?) {
    when (state) {
        1 -> {
            val c = taskC()
            continuation?.resumeWith(c)
        }

        2 -> {
            val d = taskD()
            continuation?.resumeWith(d)
        }
    }
}

fun continuationWayUnderTheHood() {

    // STATE 1 → A
    stateMachinePart1(1, object : DemoContinuation {
        override fun resumeWith(result: String) {
            Log.d(CONTINUATION_TAG, "resumeWith: $result") // A

            // STATE 2 → C
            stateMachinePart2(1, object : DemoContinuation {
                override fun resumeWith(result: String) {
                    Log.d(CONTINUATION_TAG, "resumeWith: $result") // C

                    // STATE 3 → B
                    stateMachinePart1(2, object : DemoContinuation {
                        override fun resumeWith(result: String) {
                            Log.d(CONTINUATION_TAG, "resumeWith: $result") // B

                            // STATE 4 → D
                            stateMachinePart2(2, null)
                        }
                    })
                }
            })
        }
    })
}

flow.txt

+-------------------------------+
| continuationWayUnderTheHood() |
+-------------------------------+
               |
               v
       firstFunction(1)
               |
               v
            taskA()
               |
               v
         resumeWith("A")
               |
               v
      secondFunction(1)
               |
               v
            taskC()
               |
               v
         resumeWith("C")
               |
               v
       firstFunction(2)
               |
               v
            taskB()
               |
               v
         resumeWith("B")
               |
               v
      secondFunction(2)
               |
               v
            taskD()
               |
               v
         resumeWith("D")