# Transforming to CPS

04 Aug 2018I recently picked up Appel’s classic Compiling with Continuations and have been refreshing my continuation-fu more generally.

Continuation-passing style (CPS) itself is nothing uncommon to the
functional programmer; it simply involves writing in a manner such that
functions never return, instead passing control over to something else (a
*continuation*) to finish the job. The simplest example is just the identity
function, which in CPS looks like this:

```
id :: a -> (a -> b) -> b
id x k = k x
```

The first argument is the conventional identity function argument – the second is the continuation. I wrote a little about continuations in the context of the Giry monad, which is a somewhat unfamiliar setting, but one that follows the same principles as anything else.

In this post I just want to summarise a few useful CPS transforms and related techniques in one place.

## Manual CPS Transformation

Consider a binary tree type. We’ll keep things simple here:

```
data Tree a =
Leaf a
| Branch a (Tree a) (Tree a)
```

Calculating the depth of a tree is done very easily:

```
depth :: Tree a -> Int
depth = loop where
loop tree = case tree of
Leaf _ -> 1
Branch _ l r ->
let dl = loop l
dr = loop r
in succ (max dl dr)
```

Note however that this is not a tail-recursive function – that is, it does not end with a call to itself (instead it ends with a call to something like ‘succ . uncurry max’). This isn’t necessarily a big deal – the function is easy to read and write and everything, and certainly has fine performance characteristics in Haskell – but it is less easy to deal with for, say, an optimising compiler that may want to handle evaluation in this or that alternative way (primarily related to memory management).

One can construct a tail-recursive (depth-first) version of ‘depth’ via a manual CPS transformation. The looping function is simply augmented to take an additional continuation argument, like so:

```
depth :: Tree a -> Int
depth tree = loop tree id where
loop cons k = case cons of
Leaf _ -> k 1
Branch _ l r ->
loop l $ \dl ->
loop r $ \dr ->
k (succ (max dl dr))
```

Notice now that the ‘loop’ function terminates with a call to itself (or just passes control to a supplied continuation), and is thus tail-recursive.

Due to the presence of the continuation argument, ‘loop’ is a higher-order function. This is fine and dandy in Haskell, but there is a neat technique called defunctionalisation that allows us to avoid the jump to higher-order and makes sure things stay KILO (“keep it lower order”), which can be simpler to deal with more generally.

The idea is just to reify the continuations as abstract syntax, and then
evaluate them as one would any embedded language. Note the continuation ```
\dl
-> ..
```

, for example – the free parameters ‘r’ and ‘k’ occuring in the function
body correspond to a tree (the right subtree) and another continuation,
respectively. And in `\dr -> ..`

one has the free parameters ‘dl’ and ‘k’ –
now the depth of the left subtree, and the other continuation again. We also
have ‘id’ used on the initial call to ‘loop’. These can all be reified via the
following data type:

```
data DCont a =
DContL (Tree a) (DCont a)
| DContR Int (DCont a)
| DContId
```

Note that this is a very simple recursive type – it has a simple list-like pattern of recursion, in which each ‘level’ of a value is either a constructor, carrying both a field of some type and a recursive point, or is the ‘DContId’ constructor, which simply terminates the recursion. The reified continuations are, on a suitable level of abstraction, more or less the sequential operations to be performed in the computation. In other words: by reifying the continuations, we also reify the stack of the computation.

Now ‘depth’ can be rewritten such that its looping function is not higher-order; the cost is that another function is needed, one that lets us evaluate items (again, reified continuations) on the stack:

```
depth :: Tree a -> Int
depth tree = loop tree DContId where
loop cons k = case cons of
Leaf _ -> eval k 1
Branch _ l r -> loop l (DContL r k)
eval cons d = case cons of
DContL r k -> loop r (DContR d k)
DContR dl k -> eval k (succ (max dl d))
DContId -> d
```

The resulting function is *mutually* tail-recursive in terms of both ‘loop’
and ‘eval’, neither of which are higher-order.

One can do a little better in this particular case and reify the stack using an actual Haskell list, which simplifies evaluation somewhat – it just requires that the list elements have a type along the lines of ‘(Tree a, Int)’ rather than something like ‘Either (Tree a) Int’, which is effectively what we get from ‘DCont a’. You can see an example of this in this StackOverflow answer by Chris Taylor.

## Mechanical CPS Transformation

“Mechanical CPS transformation” might be translated as simply “compiling with continuations.” Matt Might has quite a few posts on this topic; in particular he has one very nice post on mechanical CPS conversion that summarises various transformations described in Appel, etc.

Matt describes three transformations that I think illustrate the general mechanical CPS business very well (he describes more, but they are more specialised). The first is a “naive” transformation, which is simple, but produces a lot of noisy “administrative redexes” that must be cleaned up in another pass. The second is a higher-order transformation, which makes use of the host language’s facilities for function definition and application – it produces simpler code, but some unnecessary noise still leaks through. The last is a “hybrid” transformation, which makes use of both the naive and higher-order transformations, depending on which is more appropriate.

Let’s take a look at these in Haskell. First let’s get some imports out of the way:

```
{-# LANGUAGE OverloadedStrings #-}
import Data.Monoid
import Data.Text (Text)
import qualified Data.Text as T
import Data.Unique
import qualified Text.PrettyPrint.Leijen.Text as PP
```

I’ll also make use of a simple, Racket-like ‘gensym’ function:

```
gensym :: IO Text
gensym = fmap render newUnique where
render u =
let hu = hashUnique u
in T.pack ("$v" <> show hu)
```

We’ll use a bare-bones lambda calculus as our input language. Many examples – Appel’s especially – use significantly more complex languages when illustrating CPS transforms, but I think this distracts from the meat of the topic. Lambda does just fine:

```
data Expr =
Lam Text Expr
| Var Text
| App Expr Expr
```

I want to render expressions in my input and output languages in a Lisp-like
manner. This is very easy to do using a good pretty-printing library;
here I’m using the excellent *wl-pprint-text*, and will omit the ‘Pretty’
instances in the body of my post. But I’ll link to a gist including them at
the bottom.

When performing a mechanical CPS transform, one targets both “atomic” expressions – i.e., variables and lambda abstractions – and “complex” expressions, i.e. function application. The target language is thus a combination of the ‘AExpr’ and ‘CExpr’ types:

```
data AExpr =
AVar Text
| ALam [Text] CExpr
data CExpr =
CApp AExpr [AExpr]
```

All the mechanical CPS transformations use variants on two functions going by
the cryptic names **m** and **t**. **m** is responsible for converting
atomic expressions in the input languages (i.e., variables and lambda
abstractions) into atomic expressions in the target language (an atomic CPS
expression). **t** is the actual CPS transformation; it converts an expression
in the input language into CPS, invoking a specified continuation (already in
the target language) on the result.

Let’s look at the naive transform. Here are **m** and **t**, prefixed by ‘n’
to indicate that they are naive. First, **m**:

```
nm :: Expr -> IO AExpr
nm expr = case expr of
Lam var cexpr0 -> do
k <- gensym
cexpr1 <- nt cexpr0 (AVar k)
return (ALam [var, k] cexpr1)
Var var -> return (AVar var)
App {} -> error "non-atomic expression"
```

(N.b. you almost never want to use ‘error’ in a production implementation of
*anything*. It’s trivial to wrap e.g. ‘MaybeT’ around the appropriate
functions to handle the bogus pattern match on ‘App’ totally, but I just want
to keep the types super simple here.)

The only noteworthy thing that **m** does here is in the case of a lambda
abstraction: a new abstract continuation is generated, and the body of the
abstraction is converted to CPS via **t**, such that the freshly-generated
continuation is called on the result. Remember, **m** is really just mapping
atomic expressions in the input language to atomic expressions in the target
language.

Here’s **t** for the naive transform. Remember, **t** is responsible for
converting expressions to continuation-passing style:

```
nt :: Expr -> AExpr -> IO CExpr
nt expr cont = case expr of
Lam {} -> do
aexpr <- m expr
return (CApp cont [aexpr])
Var _ -> do
aexpr <- m expr
return (CApp cont [aexpr])
App f e -> do
fs <- gensym
es <- gensym
let aexpr0 = ALam [es] (CApp (AVar fs) [AVar es, cont])
cexpr <- nt e aexpr0
let aexpr1 = ALam [fs] cexpr
nt f aexpr1
```

For both kinds of atomic expressions (lambda and variable), the expression is
converted to the target language via **m**, and then the supplied continuation
is applied to it. Very simple.

In the case of function application (a “complex”, or non-atomic expression), both the function to be applied, and the argument it is to be applied to, must be converted to CPS. This is done by generating two fresh continuations, transforming the argument, and then transforming the function. The control flow here is always handled by stitching continuations together; notice when transforming the function ‘f’ that the continuation to be applied has already handled its argument.

Next, the higher-order transform. Here are **m** and **t**:

```
hom :: Expr -> IO AExpr
hom expr = case expr of
Lam var e -> do
k <- gensym
ce <- hot e (\rv -> return (CApp (AVar k) [rv]))
return (ALam [var, k] ce)
Var n -> return (AVar n)
App {} -> error "non-atomic expression"
hot :: Expr -> (AExpr -> IO CExpr) -> IO CExpr
hot expr k = case expr of
Lam {} -> do
aexpr <- m expr
k aexpr
Var {} -> do
aexpr <- m expr
k aexpr
App f e -> do
rv <- gensym
xformed <- k (AVar rv)
let cont = ALam [rv] xformed
cexpr fs = hot e (\es -> return (CApp fs [es, cont]))
hot f cexpr
```

Both of these have the same form as they do in the naive transform – the difference here is simply that the continuation to be applied to a transformed expression is expressed in the host language – i.e., here, Haskell. Thus the transform is “higher-order,” in exactly the same sense that higher-order abstract syntax is higher-order.

The final transformation I’ll illustrate here, the hybrid transform, applies
the naive transformation to lambda and variable expressions, and applies the
higher-order transformation to function applications. Here **t** is split up
into **tc** and **tk** to handle these cases accordingly:

```
m :: Expr -> IO AExpr
m expr = case expr of
Lam var cexpr -> do
k <- gensym
xformed <- tc cexpr (AVar k)
return (ALam [var, k] xformed)
Var n -> return (AVar n)
App {} -> error "non-atomic expression"
tc :: Expr -> AExpr -> IO CExpr
tc expr c = case expr of
Lam {} -> do
aexpr <- m expr
return (CApp c [aexpr])
Var _ -> do
aexpr <- m expr
return (CApp c [aexpr])
App f e -> do
let cexpr fs = tk e (\es -> return (CApp fs [es, c]))
tk f cexpr
tk :: Expr -> (AExpr -> IO CExpr) -> IO CExpr
tk expr k = case expr of
Lam {} -> do
aexpr <- m expr
k aexpr
Var {} -> do
aexpr <- m expr
k aexpr
App f e -> do
rv <- gensym
xformed <- k (AVar rv)
let cont = ALam [rv] xformed
cexpr fs = tk e (\es -> return (CApp fs [es, cont]))
tk f cexpr
```

Matt illustrates these transformations on a simple expression: `(g a)`

. We can
do the same:

```
test :: Expr
test = App (Var "g") (Var "a")
```

First, the naive transform. Note all the noisy administrative redexes that come along with it:

```
> cexpr <- nt test (AVar "halt")
> PP.pretty cexpr
((λ ($v1).
((λ ($v2).
($v1 $v2 halt)) a)) g)
```

The higher-order transform does better, containing only one such redex (an eta-expansion). Note that since the supplied continuation must be expressed in terms of a Haskell function, we need to write it in a more HOAS-y style:

```
> cexpr <- hot test (\ans -> return (CApp (AVar "halt") [ans]))
> PP.pretty cexpr
(g a (λ ($v3).
(halt $v3)))
```

Finally the hybrid transform, which, here, is literally perfect. We don’t even need to deal with the minor annoyance of the HOAS-style continuation when calling it:

```
> cexpr <- tc test (AVar "halt")
> PP.pretty cexpr
(g a halt)
```

Matt goes on to describe a “partioned CPS transform” that can be used to recover a stack, in (seemingly) much the same manner that the defunctionalised manual CPS transform worked in the previous section. Very neat, but something I’ll have to look at in another post.

## Fin

CPS is pretty gnarly. My experience in *compiling* with continuations is not
substantial, but I dig learning it. Appel’s book, in particular, is meaty –
expect more posts on the subject here eventually, probably.

‘Til next time! I’ve dumped the code from the latter section into a gist.