🚀 · Yield vs. Return & How Pausing Functions Work

8 min read · Updated on by

Read on for an introduction to the sequence method, an explanation of yield vs. return, a look into how to implement functions that can pause using continuations, and how this relates to suspend functions and coroutines.

sequence


//sampleStart
fun <T> sequence(
    block: suspend SequenceScope<T>.() -> Unit
): Sequence<T>
//sampleEnd

The sequence method is a unique way you can build up a Sequence by sequentially generating each element (or collection of elements). Even from it’s signature, you can see it’s a different — we’ve never encountered the suspend keyword until now.

The key difference between generateSequence, which we discussed in the previous article, and sequence is that in this case, the values are not produced by returning them from a lambda, but by calling yield (for producing a single element), or yieldAll (for producing all elements in a collection). Crucially, neither of those calls need to be at the end of a function, like a return statement — you can call them whenever you want to, which allows us to solve certain problems much more elegantly.

Let’s take a look at an example:


//sampleStart
fun fibonacci_return(): Sequence<Int> {
    var terms = Pair(0, 1)
    
    return generateSequence {
        val returnValue = terms.first
        terms = Pair(terms.second, terms.first + terms.second)
        returnValue
    }
}

fun fibonacci_yield() = sequence {
    var terms = Pair(0, 1)
    
    while (true) {
        yield(terms.first)
        terms = Pair(terms.second, terms.first + terms.second)
    }
}
//sampleEnd
fun main() {
    val poem = """
        Kotlin, the architect of code's tower,
        With sealed classes, it builds the power.
        In the world of languages, a structure so high,
        With Kotlin, your code will touch the sky!
    """.trimIndent()
    println(poem)
}

Both of the above implement the sequence of Fibonacci numbers.

In the first version, where we’re constrained to returning from a function, it’s more difficult to deal with state (notice that we have to keep the state outside of the lambda defining the sequence), more difficult to update state after we produce the value (notice that we need to save the value in an intermediate variable so we can return it) and involves other nuisances as well. For example, since the state is kept track of separately, there’s no input to the lambda in generateSequence. However, this means that we’re using the generateSequence variant that produces Sequences constrained to run only once! In a way, this is actually a good thing, since if we ran it multiple times, we would get nonsensical results — both runs would be accessing the same terms variable — but it’s certainly not something we want to deal with in a scenario as simple as generating a mathematical sequence of numbers.

The yielding version solves all of these problems nicely. The fundamental reason this is so is because it gives us complete control over the flow of execution.

Yielding vs. Returning

When using constructs like generateSequence, the order in which certain things happen are prescribed, and cannot be changed. For instance, in generateSequence, the prescribed flow of execution is:

  1. run function,
  2. use produced value as element in sequence,
  3. pass produced value as input to next run of function,
  4. repeat.

These steps are set in stone, and it is impossible to e.g. wedge an additional step between “use produced value as element in sequence” and “pass produced value as input to next run of function”, which is ideally what we want to do when generating a Fibonacci sequence. We also don’t have control over the manner in which these steps are repeated — there is obviously some sort of loop or iteration involved, but it’s hidden inside the implementation of generateSequence.

When we produce values by calling a function (commonly called yielding) instead of returning, we can do whatever we want, including updating state after producing a value (which saves us from having to define a temporary variable), and implementing the loop manually (which allows us to include the state inside the generating lambda, but still keep it outside the loop, which is why we couldn’t do the same thing with generateSequence).

Now, if you’ve been paying attention, you should be thinking Waaait a second…how can that even work? After all, the defining feature of sequences is that all operations get applied as soon as an element is produced — indeed, sequences can be infinite. But in the approach above, we’re not producing an element at the end of the execution of the lambda, but in the middle of it. In essence, this amounts to being able to instruct the function to produce a value and pause, and having it continue from the same spot when the next value is requested.

That sounds crazy — how could that possibly work? How can a function just decide to pause (or temporarily suspend its execution) in the middle of a run, and resume later? And what is that suspend thing?

The answer to the second question is: actually, it doesn’t really pause. It just looks like does, which is a trick made possible by the Kotlin compiler.

Implementing a pausing function

Let’s try to demonstrate the essence of how to implement a “pausing” function without taking advantage of Kotlin-specific compiler-fu, so you get a very rough idea of what the compiler is doing behind the scenes. The actual technical implementation uses the framework upon which Kotlin coroutines are built, which is also where that suspend keyword comes from.


//sampleStart
sealed interface PauseState
object ValueReady : PauseState
object ValueNotReady : PauseState
object Finished : PauseState

typealias SimpleContinuation = () -> Unit

class PausingIterator<T : Any>(firstPart: PausingIterator<T>.() -> Unit): Iterator<T> {
    private var state: PauseState = ValueNotReady
    private var value: T? = null
    private var nextPart: SimpleContinuation? = { firstPart() }

    fun yield(value: T, nextPart: SimpleContinuation) {
        this.value = value
        this.nextPart = nextPart
        state = ValueReady
    }

    override fun hasNext(): Boolean {
        while(true) {
            when(this.state) {
                ValueReady -> return true
                Finished -> return false
                ValueNotReady -> { /* Do nothing */}
            }
            this.state = Finished
            nextPart!!()
        }
    }

    override fun next() = when(state) {
        Finished -> throw NoSuchElementException()
        ValueReady -> {
            state = ValueNotReady
            value!!
        }
        ValueNotReady -> {
            if(!hasNext()) throw NoSuchElementException()
            value!!
        }
    }
}

fun <T: Any> buildUsingYield(
    block: PausingIterator<T>.() -> Unit
): Iterator<T> = PausingIterator<T>(block)

fun main() {
    val seq_rolledOut = buildUsingYield<Int> {
        var state = 1
        println("Yielding 1 and pausing")
        yield(state++) {
            println("Resumed after 1, yielding 2 and pausing")
            yield(state++) {
                println("Resumed after 2, yielding 3 and pausing")
                yield(state++) {
                    println("Resumed after 3, and finishing")
                }
            }
        }
    }
    
    val seq_withLooping = buildUsingYield<Int> {
        var state = 1
        fun buildContinuation(): SimpleContinuation = {
            if(state <= 3) {
                println("Resumed after ${state-1}, yielding ${state} and pausing")
                yield(state++, buildContinuation())
            } else {
                println("Resumed after ${state-1} and finishing")
            }
        }


        println("Yielding 1 and pausing")
        yield(state++, buildContinuation())
    }

    seq_rolledOut.forEach { 
        println(it)
        println("Yadayada")
    }
    
    println("--------")
    
    seq_withLooping.forEach { 
        println(it)
        println("Yadayada")
    }
}
//sampleEnd

The above is a simplified builder for Iterators producing values of some non-nullable type T via a yield “pausing” function. It uses a simple state machine to determine if a new value has already been produced and is available, or is not yet available and must be produced. To keep things as simple as possible, we’re completely ignoring any type of error handling except for the basic “no more elements” scenario, and also don’t support producing nullable values.

The fundamental principle that makes this work is the fact that yield accepts a callback, which contains everything that should happen after the function is “resumed”. When yield is called, this callback is saved in nextPart, and gets called once the state machine determines that we need to produce a new value. These types of callbacks are called continuations (they continue with executing the rest of the code), and they allow for some of the more powerful techniques in programming.

Once you think about it like that, you’ll see that there isn’t actually any pausing involved — it’s a trick. Instead, you break up your function into nested callbacks (or, to use a more technical term, rewrite it in continuation-passing style) and pass each callback to some sort of logic that decides when to call them. It is then trivial to call each successive callback whenever you feel like it — now, later, whenever.

In the above, the state machine controls when the function gets “resumed” (it gets resumed when hasNext is called), but you could do anything you want. The actual implementation of sequence uses something similar, but much more robust (deals with errors, allows nullable values, and so on).

“But we’re not writing any such code when building sequences!” I hear you exclaim. “There are no callbacks, no nested lambdas, nothing like that!”. That’s because the compiler is doing all this for us, behind the scenes. To simplify things greatly, that suspend keyword you saw at the beginning tells the compiler to take the following:


//sampleStart
val seq = buildUsingYield<Int> {
  println("Yielding 1 and pausing")
  yield(1)
  println("Resumed after 1, yielding 2 and pausing")
  yield(2)
  println("Resumed after 2, yielding 3 and pausing")
  yield(3)
  println("Resumed after 3, and finishing")
}
//sampleEnd

And transform it by adding a hidden Continuation parameter to yield, and change the calling code to something that distantly resembles the following:


//sampleStart
val seq = buildUsingYield<Int> {
    println("Yielding 1 and pausing")
    yield(1) {
        println("Resumed after 1, yielding 2 and pausing")
        yield(2) {
            println("Resumed after 2, yielding 3 and pausing")
            yield(3) {
                println("Resumed after 3, and finishing")
            }
        }
    }
}
//sampleEnd

Just so we understand each other, theContinuation instance is not a functional or SAM type— you can’t pass a lambda where a Continuation is expected. In the above, I’m simply demonstrating what the Continuation represents — a thing you can call to continue executing the rest of the program — by writing a lambda instead of it. You can see that that can’t be exactly what’s happening, because you couldn’t translate our Fibonacci example, where yield is part of a while loop, into a similar form. However in principle, this — i.e. rewriting yield to accept a Continuation— is what happens.

By the way, this is the fundamental principle the entire coroutines framework is built upon, but that’s a story for another day.

It’s worth mentioning that there is a builder for instances of Iterator that permits the same syntax:


//sampleStart
fun <T> iterator(
    block: suspend SequenceScope<T>.() -> Unit
): Iterator<T>
//sampleEnd

Example


//sampleStart
interface Node
data class Tree(val node: Node, val children: List): Iterable {
    override fun iterator(): Iterator = iterator {
        yield(node)
        children.forEach { yieldAll(it) }
    }
}
//sampleEnd
fun main() {
    val poem = """
        In the code's carnival, Kotlin's the ride,
        With extension functions, it's the guide.
        From loops to spins, a coding spree,
        In the world of development, it's the key!
    """.trimIndent()
    println(poem)
}
Previous🚗 · How They Work & How To Use Them
Next 🚗 · Iterable vs. Sequence vs. Java Stream

Leave a Comment

This site uses Akismet to reduce spam. Learn how your comment data is processed.

The Kotlin Primer