A Compiler for Multi-shot Effect Handlers

The pond at Pembroke College, Cambridge. I have seen deeper rain puddles in Edinburgh.

I recently spent two weeks at OCaml Labs, the University of Cambridge implementing a compiler for the experimental, research language Links with effect handlers. During those two weeks I managed to compile a substantial subset of the Links language to native code, in particular, I managed to compile some programs that use handlers. My Links compiler reuses most of the infrastructure from the OCaml effects/multicore compiler.

At OCaml Labs I was working with KC Sivaramakrishnan, whom is working on the Multicore OCaml project.

In this post I shall discuss numerous three things: I will briefly describe the compiler infrastructure, discuss an example program that compiles, and how I encode multi-shot handlers.

Compiler infrastructure

My stay at OCaml Labs was really productive. I managed to implement a compiler for an interesting subset of Links with effect handlers. It is still very much work in progress as there are many Links primitives that the compiler does not yet support. The source code for the compiler is available at the Links repository in the effect-handlers-compilation branch.

The reason I was able to implement a compiler for effect handlers this fast is that I am hooking into the OCaml effects/multicore compiler backend. The OCaml multicore project uses effect handlers to provide multicore support. A consequence of reusing the OCaml infrastructure is that I get the translation of Links into native code almost for free. Furthermore, I had direct access to people who are familiar with the OCaml compiler infrastructure. This made the process of integrating Links with infrastructure smooth.

The OCaml backend embodies several intermediate representations (IRs). Figure 1 depicts a diagram representation of the OCaml backend.

Figure 1. A diagrammatic view of the compiler infrastructure.

There are several potential entry points, but for maximum flexibility I chose to translate the Links IR into the OCaml Lambda IR. This gives the option of compiling to either byte or native code. At the moment I am only focusing on native branch. But in the future it may be interesting to take advantage of compilers like js_of_ocaml that translates OCaml byte code into JavaScript.

In the most recent OCaml compiler Lambda gets translated into Flambda, which is an optimisation framework. But since Links does not compile with the latest OCaml compiler (due to third party dependencies) I am bypassing this phase for now. However, I plan to merge Flambda back in at a later point.

Coin tossing example running natively

I have previously used the coin tossing example to demonstrate programming with handlers. It is the smallest, non-trivial example that demonstrates the handler abstraction that I am familiar. The simplicity of the example makes it a good first program to compile.

I shall briefly summarise the coin toss example. The example comprise a nondeterministic choice operation and an exception operation, which we name Choose and Fail respectively:

## Non-deterministic choice operation
sig choose : () {Choose:Bool |_}~> Bool
fun choose() { do Choose }

## Failure (exception) operation
sig fail : () {Fail:Zero |_}~> _
fun fail() { switch(do Fail) { } }

Here Zero is the empty type. The empty case statement switch (...) { } conveniently lets us assign a polymorphic type to the function fail such that we may use it in any context. We use the operations to model a coin toss and a drunk coin toss:

## Function modelling a coin toss
sig toss : () {Choose:Bool|_}~> Toss
fun toss() {
  if (choose()) Heads
  else Tails
}

## Function modelling a drunk toss
sig drunkToss : () {Choose:Bool,Fail:Zero |_}~> Toss
fun drunkToss() {
  if (choose()) toss()
  else fail()
}

After flipping the coin a drunkard may fail to catch it. We capture this in drunkToss by nondeterministically choosing between toss and fail. In the example, we have two different interpretations of Choose:

• randomResult randomly interprets the operation as true or false.
• allResults enumerates the sample space by trying both truth values.

The latter interpretation has the most interesting implementation:

sig allResults : (Comp({Choose:Bool |e}, a)) -> Comp({Choose{_} |e}, [a])
fun allResults(m)() {
  open handle(m) {
     case Return(x) -> [x]
     case Choose(k) -> k(true) ++ k(false)
  }
}

It enumerates the sample space by first interpreting Choose as true and afterwards as false. The result of an interpretation is a list of outcomes. The Return-clause lifts the outcome of a coin flip into a singleton list. In the Choose-clause the outcomes of either interpretation get concatenated. The open handle is the effect forwarding construct in Links. An aside: I am actually planning to get rid of this construct such that I have only one construct handle that will forward by default.

For the operation Fail we also have two interpretations:

• maybeResult that returns Nothing when Fail occurs.
• persevere that recursively restarts the computation when Fail occurs.

Both handlers are simple to implement, here is maybeResult:

sig maybeResult : (Comp({Fail:Zero |e}, a)) -> Comp({Fail{_} |e}, Maybe(a))
fun maybeResult(m)() {
  open handle(m) {
    case Return(x) -> Just(x)
    case Fail(_)   -> Nothing
  }
}

In the Return-clause we tag the outcome of a coin flip with Just. But if the drunkard fails to catch the coin again, then in the Fail-clause we return Nothing.

I have put together the entire source code in the file cointoss-native.links. Currently, the Links compiler can translate the entire example to native code. The listing below shows a possible output of running the binary cointoss produced by the Links compiler:

$ ./cointoss
# examples
#1 randomResult(toss):
Heads

#2 allResults(toss):
[Heads,Tails]

#3 maybeResult(randomResult(toss)):
Just(Heads)

#4 allResults(maybeResult(toss)):
[Just(Heads),Just(Tails)]

#5 maybeResult(allResults(toss)):
Just([Heads,Tails])

#6 maybeResult(randomResult(drunkToss)):
Nothing

#7 persevere(randomResult(drunkToss)):
Tails

#8 maybeResult(allResults(drunkToss)):
Nothing

#9 allResults(maybeResult(drunkToss)):
[Just(Heads),Just(Tails),Nothing]

Notice that there is an obvious, possible interpretation missing from the listing, namely: allResults(persevere(drunkToss)). This interpretation causes the program to diverge as allResults stubbornly keeps picking the false-branch in drunkToss.

Encoding multi-shot handlers

The OCaml language supports only affine handlers, i.e. handlers that consume their continuations at most once. The handler allResults is an example of a multi-shot handler, it consumes its continuation twice: k(true) and k(false).

Before going to OCaml Labs I was wondering how to encode multi-shot handlers in OCaml Lambda. To my surprise it turned out to be really easy. A few months ago KC implemented a Obj.clone which clones a given object. We can use this to simulate multi-shot handlers by cloning the continuation before consuming it. Consider part of the coin tossing example in OCaml:

type toss = Heads | Tails

effect Choose : bool

let choose () = perform Choose

let toss () = if (choose ())
              then Heads
              else Tails

let all_results m =
  match m () with
  | v -> [v]
  | effect Choose k -> let k' = Obj.clone k in
                       (continue k true) @ (continue k' false)

The function continue takes a continuation and the parameter that the continuation should be applied to. Aside: I think OCaml has a neat uniform syntax for exception, effect and value handling. The former two are simply syntactic extensions to the match ... with (or try ... with) construct.

Running the above program in an OCaml REPL, we obtain the desired result:

# all_results toss;;
- : toss list = [Heads; Tails]

Leaky abstraction

Although, we can simulate multi-shot handlers using Obj.clone it provides a fragile abstraction. To see this consider a silly extension to the coin tossing example. We introduce a linear handler chooseNot that invokes and returns the negation of Choose:

let choose_not m =
  match m () with
  | v -> v
  | effect Choose k -> let choice = choose () in
                       continue k (not choice)

If we compose all_results with choose_not to interpret toss, we get a run-time error:

# all_results (fun () -> choose_not toss);;
Exception: Invalid_argument "continuation already taken".
Called from unknown location

OCaml complains that we consume the linear continuation in choose_not twice. It may not be immediately clear how we ended up consuming it two times. A continuation is essentially a pointer to the address from where computation should be resumed. The problem is that Obj.clone shallowly clones the call stack, thus k and k' in all_results point to the exact same continuation in choose_not. In the presence of a multi-shot handler, every linear handler must be promoted¹ to a multi-shot handler.

My implementation in Links repairs this deficiency. The same program runs without errors in Links. The following is the implementation in Links:

sig chooseNot : (Comp({Choose:Bool |e}, a)) -> Comp({Choose:Bool |e}, a)
fun chooseNot(m)() {
  open handle(m) {
     case Return(x) -> x
     case Choose(k) -> k(not(choose()))
  }
}

I include this in the source file as the tenth example, it outputs:

#10 allResults(chooseNot(toss)):
[Tails,Heads]

As we would expect it returns the reverse of allResults(toss). My solution is rather simple, and overly conservative: I encode every handler as a multi-shot handler by always copying the continuation before invocation. In plain OCaml my encoding amounts to:

| effect Choose k ->
   let fake_k = fun param ->
     let k' = Obj.clone k in
     continue k' param
   in
   (fake_k true) @ (fake_k false)

The problem with this approach is that we lose the performance benefit of linear handlers. But fixing this problem is one of the next phases in my project.

Conclusion

I am grateful to OCaml Labs for letting me visit them and providing me with much appreciated help.

To the best of my knowledge I am the first to compile the full abstraction of (deep) handlers to native code. But there is much work that remains to be done. Furthermore, in this post I have not provided any performance numbers. I soon start measuring and I intend to post some of the early results here. The next step is to implement optimisations such that I can recover the performance benefits of linear handlers.

For your convenience I provide the full source code for the examples:

• cointoss-native.links (runs with Links git commit #b477b74)
• coin_example_leaky_abstraction.ml (runs with OCaml-4.02.2+effects)