Better effect type management with coroutines

December 14, 2022

One interesting language design thread we've stumbled upon with Swift, but haven't fully developed, is the use of coroutines as an alternative to higher-order functions for building "control flow like" constructs. No slight against higher-order functions—as a singular language concept, they definitely pay for themselves in the expressivity they add—but in typed languages, whose type systems already have a tendency toward ever-growing complexity, function types are a particularly good complexity generator. Coroutines are definitely a more complex concept than functions, but I wanted to share some of my thoughts on how they have the potential to save complexity elsewhere, allowing for similar expressivity in extensible control flow constructs while requiring a less complex type system to manage the composition of effects within those constructs.

The effect of effects on function types

One common pattern with higher-order functions is resource management. We can define a function that sets up a resource, executes a body function provided as a parameter, then tears down the resource when the body finishes running (in a nonspecific mostly-imperative mainstream language syntax):

func withResource(body: (Resource) -> ()) {
  var r = setUpResource()
  body(r)
  tearDownResource(r)
}

withResource {|r|
  enjoyResource()
}

This is all nice and simple, and languages these days like to spend a lot of syntax sugar on making higher-order functions like this feel lightweight and as "native" as primitive control flow constructs like if and for. However, the illusion tends to fall apart once we consider the kinds of effects we want the body function to be able to have. Maybe the language allows for flow-sensitive initialization of variables, raising an error when a variable is used before it's initialized:

var magicNumber

print(magicNumber) // Error, magicNumber not initialized

if condition {
  magicNumber = 69
  print(magicNumber) // We can use magicNumber here because it's initialized
}

print(magicNumber) // Error, magicNumber not initialized when `condition` is false

It's reasonable to expect that this would work in a withResource block too. Looking at its type signature alone, a type checker can tell that withResource takes a parameter of function type, but doesn't know what the implementation does with it—it could execute it once, or twice, or any number of times, or ignore the parameter altogether and not execute it at all, or save it away to call later. So, without more information in the type signature for withResource, we can't tell whether it's statically sound to let the closure capture an uninitialized variable, nor whether we can assume the closure was executed

var magicNumber

withResource {|r|
  magicNumber = 69 // Error, can this closure capture `magicNumber` uninitialized?
  print(magicNumber)
}

print(magicNumber) // Error, did `withResource` actually execute the closure?

So let's make up a @runExactlyOnce modifier that withResource can use to promise that it runs the body exactly once:

func withResource(body: @runExactlyOnce (Resource) -> ()) {
  var r = setUpResource()
  body(r)
  tearDownResource()
}

var magicNumber

withResource {|r|
  magicNumber = 69 // OK to capture magicNumber uninitialized, we know this'll run once
  print(magicNumber)
}

print(magicNumber) // OK, we know the closure ran once

That's not so bad, but maybe we also want to be able to return a value from the body through the resource management, so we can use withResource in an expression. Let's say the language already has parametric types/generics, so we can say:

func withResource<T>(body: @runExactlyOnce (Resource) -> T) -> T {
  var r = setUpResource()
  var result = body(r)
  tearDownResource(r)
  return result
}

var magicNumber = withResource {|r| getMagicNumberFromResource(r) }

But maybe the language also has exceptions, and uses the type system to distinguish functions that can throw exceptions from ones that can't. withResource ought to work as a way to make sure that the resource gets torn down even if an uncaught exception occurs. But it should still be treated as a function that doesn't throw an error itself when the body function doesn't. So we need to make up an annotation now for "function parameter that may throw an exception" and "function that throws an exception when its parameter does":

func withResource<T>(body: @runExactlyOnce () throws -> T) rethrows -> T {
  var r = setUpResource()
  try {
    var result = body(r)
    tearDownResource(r)
    return result
  } catch error {
    tearDownResource(r)
    throw error
  }
}

var magicNumber = withResource {|r| tryToGetMagicNumberFromResource(r) }

Extensible effect systems are pretty great, though, and we could anticipate instead that there will be more kinds of interesting effects in the language in the future, and introduce a new kind of generic effect parameter instead:

func withResource<T, E: Effect>(body: @runExactlyOnce (Resource) E -> T) E -> T { ... }

var magicNumber = withResource {|r| tryToGetMagicNumberFromResource(r) }

At each step, as we try to expand the expressivity of the base language, it takes more and more effort for a "control-flow-like" higher-order function to maintain the illusion of being a control flow construct.

Detangling the caller and callee contexts

When we build up these higher-order functions, they work by executing their function parameters inside their own context, and the type system complexity arises from having to communicate "if executing the function parameter does this, then so does executing the higher-order function as whole". Also, to feel like control flow, we want people using these constructs to think of the function argument as being part of the surrounding code, but we're really capturing some parts of the surrounding code and wrapping them up into an independent function, a third context distinct from either the caller code or the callee, and we also grow type system complexity by trying to leash the capabilities of that third context to not wander too far from either of the two primary contexts we care about, with the @runExactlyOnce modifier.

Instead of looking at withResource as a higher-order function, we could look at it as a coroutine that interleaves its execution with the caller: instead of the caller passing a function for it to execute, withResource could yield control back to the caller with the resource, and then tear down the resource when the caller resumes it:

coroutine resource: () -> (Resource) -> () {
  var r = setUpResource()
  yield r
  tearDownResource(r)
}

We could specify some kind of control flow construct to consume a single-shot coroutine like this, maybe something like C♯'s using blocks:

var magicNumber
using r = resource() {
  magicNumber = tryToGetMagicNumberFromResource(r)
}
print(magicNumber)

Since the body of the using construct is still formally part of the caller, its effects don't need to be the coroutine's business; to describe the type of the coroutine, we need to specify the values that we need to put in and expect to get out at each step of the coroutine's execution, but it's not carrying a function value to execute on the caller's behalf anymore, so it no longer needs a bunch of polymorphic variables to describe all the possible variations in the body's effects and how they project onto the coroutine's execution.

As another example, let's consider the eternal "forEach method vs. for loop" debate. It's appealing to express iteration through a collection as a higher-order function, instead of a bespoke control flow construct:

var sum = 0
forEach(magicNumbers) {|number|
  sum += number
}

Expressing the type of such a higher-order function tends to be similarly intricate to the with case, since we need to account for all the possible effects of the body, allow for early termination of the loop, and so on:

// We don't have the benefit of a builtin "break" or "continue" statement to say
// whether to keep going after each iteration
enum KeepGoing { continue, break }

func forEach<Element, E: Effect>(
  collection: Collection<Element>,
  body: @runsZeroOneOrManyTimes (Element) E -> KeepGoing
) -> () {
  var index = collection.startIndex, endIndex = collection.endIndex
  while index < endIndex {
    var keepGoing = body(collection[index])
    if keepGoing == .break {
      return
    }
    ++index
  }
}

However, writing the body of a forEach function is usually more natural than implementing some stepwise iteration protocol that can be pumped by a bespoke for loop. Many languages have converged on a coroutine-based alternative to this problem, in the form of generators:

// Borrowing a notion from regex syntax, let's say -> (Type)* means
// that a value of `Type` gets yielded zero, one, or many times
coroutine each<Element>:
  (collection: Collection<Element>) -> (Element)* -> () {
  var index = collection.startIndex, endIndex = collection.endIndex
  while index < endIndex {
    yield collection[index]
  }  
}

which a for loop can consume:

var sum = 0
for number in each(magicNumbers) {
  sum += number
}

Again, by treating the for loop body as still being part of the caller context, instead of a separate function value for the callee to invoke on the caller's behalf, the coroutine type can stay concerned with the coroutine's own business of what parameters it takes in and what values it yields, and doesn't need a bunch of extra polymorphism and annotation to pass through a body function's possible effects.

There's still no free lunch

To be clear, this isn't necessarily a net complexity win. Coroutines provide a nice way of avoiding type system complexity for a specific problem, but introduce their own complexities. Functions are relatively simple things that you pass values into and get values back, whereas coroutines can, in their full generality, accept, yield, and return progressively different kinds of values as they execute, possibly even going down different paths depending on their parameters and program state. I don't have a good sense of what a "complete" type system for coroutines would look like.

In this post I've also largely handwaved what it looks like to launch and consume coroutines on the caller's side: "use a using block for resource-like coroutines, and a for loop for generator-like coroutines" satisfies simple cases, but again in full generality, it's not great if every shape of coroutine needs a bespoke language construct to consume. You probably also want to be able to step through a coroutine manually, or pass a handle around to a partially-executed coroutine to process in different parts of the codebase. This probably requires some sort of Rust-like affine type system to model within a type system if coroutines can have potentially different inputs and yields at different points of their execution.

That all said, I still think this would be an interesting alternative language design path to explore. Higher-order functions are one of those "good enough" concepts that, while definitely great and deserving of their wide adoption, also reduce interest in exploring alternatives to the problems they solve.

View and comment on this post in the fediverse, or email me about it.