Continuation Passing Style

Quick notes on basic CPS conversion after struggling to understand ¹ and finally getting it thanks to ² (+countless others).

1. Lambda-calculus
2. CPS (call-by-value)
   1. Properly tail-recursive CPS
3. Extensions
   1. Conditionals
4. Partitioned CPS
5. References

Lambda-calculus

We start with the pure lambda calculus syntax with variables, abstractions and applications:

which we implement in Haskell as:

data Expr = Var String | Abs String Expr | App Expr Expr

CPS (call-by-value)

Converting from the Direct Style (DS) call-by-value lambda calculus to Continuation-Passing-Style (CPS) is pretty straightforward: each converted term is an abstraction that accepts an extra continuation parameter \(\kappa\) representing the rest of the program, and applies it where needed:

This conversion procedure can be given the following type:

cps :: Expr -> Expr

However, we will need some bookkeeping for generating fresh variables during the conversion. For this reason, our conversion will require a State monad:

-- cps state monad
import Control.Monad.State

type CPS a = State Int a

-- generate unique name
gensym :: String -> CPS String
gensym prefix = do
  counter <- get
  put (counter + 1)
  return (prefix ++ (show counter))

With that out of the picture, a straightforward adaptation of the above conversion gives:

-- continuation-passing-style conversion (naive)
cps :: Expr -> CPS Expr

-- var
cps (Var name) = do
  k <- gensym "k"
  return (Abs k (App (Var k) (Var name)))

-- abs
cps (Abs arg body) = do
  k <- gensym "k"
  b <- cps body
  return (Abs k (Abs arg b))

-- app
cps (App func arg) = do
  k <- gensym "k"
  f <- gensym "f"
  a <- gensym "a"
  func <- cps func
  arg <- cps arg
  return (Abs k
              (App func (Abs f
                             (App arg (Abs a
                                           (App (App (Var f) (Var a)) (Var k)))))))

Unfortunately, this naive conversion introduces quite a lot of so-called administrative redexes. A first thing to notice is that all cases produce an abstraction (dynamic abstractions), so we can try to lift these to compile-time abstractions (static abstractions), and apply them during translation so they won’t show up in the result. For instance, when encountering \(\app{\cps{x}}{k}\), we would like to \(\beta\)-reduce \(\app{\lambda\kappa.\app{\kappa}{x}}{k}\) to \(\app{k}{x}\) during translation. Our procedure becomes:

cps :: Expr -> Expr -> CPS Expr

cps (Var name) k = return (App k (Var name))

cps (Abs arg body) k = do
  c <- gensym "c"
  body <- cps body (Var c)
  return (App k (Abs arg (Abs c body)))

cps (App func arg) k = do
  f <- gensym "f"
  a <- gensym "a"
  rest <- cps arg (Abs a (App (App (Var f) (Var a)) k))
  cps func (Abs f rest)

The situation is improved, but we can do better: since we know that \(\kappa\) is always a continuation, we can pass a static continuation instead. The conversion function type thus becomes:

cps :: Expr -> (Expr -> CPS Expr) -> CPS Expr

Converting variables is straightforward: we just need to apply the static continuation to our variable:

cps (Var name) k = k (Var name)

For the rest, we mostly need to keep the type-checker happy. When converting abstractions, we bump into the following issue:

cps (Abs arg body) k = k (Abs arg (cps body ????))

The static continuation k is applied to the converted abstraction, but we still need a (static) continuation to convert the function body. The converted function body will pass its result r to whatever (dynamic) continuation gets passed to the converted function on runtime, so we simply give it a name c and wrap its application in a static continuation:

cps (Abs arg body) k = do
  c <- gensym "c"
  body <- cps body (\r -> return (App (Var c) r))
  k (Abs arg (Abs c body))

Similarly, when converting applications, we need to apply our static continuation k to the application result. So again, we name this result r and construct a dynamic continuation that will apply k to it:

cps (App func arg) k = do
  r <- gensym "r"
  rest <- k (Var r)
  cps func (\f -> cps arg (\a -> return (App (App f a) (Abs r rest))))

We just obtained what is known as the single-pass, higher-order cps conversion.

Properly tail-recursive CPS

In the case of a tail-call (that is, when converting an abstraction whose body is an application), we see that the conversion first introduces a static continuation wrapping a named dynamic continuation (when converting the abstraction), only to wrap it a second time into a dynamic continuation (when converting the application).

Instead of wrapping twice, we could simply pass along the named dynamic continuation k when converting the application:

cps_tail :: Expr -> Expr -> CPS Expr
cps_tail (App func arg) k = 
  cps func (\f -> cps arg (\a -> return (App (App f a) k)))

and have cps handle this special case when converting abstractions:

cps (Abs arg body) k = do
  c <- gensym "c"
  body <- cps_tail body (Var c)
  k (Abs arg (Abs c body))

The rest is pretty straightforward:

cps_tail (Var name) k = return (App k (Var name))

cps_tail (Abs arg body) k = do
  c <- gensym "c"
  body <- cps_tail body (Var c)
  return (App k (Abs arg (Abs c body)))

Extensions

Let us add a few constructs to our language:

And its implementation:

-- lambda calculus + conditionals
data Expr
  = Var String
  | Abs String Expr
  | App Expr Expr
  | Cond Expr Expr Expr

Conditionals

Conditionals are straightforward: we convert the condition, test its value, and convert branches continuing with our initial continuation:

The implementation follows closely, but we do not want to duplicate the dynamic continuation (which may be large) in both branches of the test. At least, we do not want to duplicate it yet. So we give it a name:

cps_tail (Cond pred conseq alt) k =
  cps pred (\p -> do
               c <- gensym "c"
               conseq <- cps_tail conseq (Var c)
               alt <- cps_tail alt (Var c)
               return (App (Abs c (Cond p conseq alt)) k))

For cps, we simply delegate to cps_tail by wrapping our static continuation behind a dynamic continuation, just like we did earlier for applications:

-- wrap static continuation into dynamic continuation
wrap :: (Expr -> CPS Expr) -> CPS Expr
wrap k = do
  r <- gensym "r"
  rest <- k (Var r)
  return (Abs r rest)

cps (Cond pred conseq alt) k =
  cps pred (\p -> do
               k <- wrap k
               conseq <- cps_tail conseq k
               alt <- cps_tail alt k
               return (Cond p conseq alt))

In terms of wrap, converting applications becomes simply:

cps (App func arg) k = do
  k <- wrap k
  cps_tail (App func arg) k

Partitioned CPS

In CPS form all calls are tail calls, which can be implemented as jumps: we can execute CPS programs without a stack. However, we may want to keep a stack e.g. for separate compilation. In this case, it is often useful to separate CPS terms originating from source abstractions (for which we may want to keep a call stack) and the ones introduced by the conversion, representing control flow (which we may implement as jumps). Symmetrically, applications can be partitioned into function calls and continuation calls.

References

e-mail comments/suggestions