Continuous-Passing Conversion in Haskell

http://en.wikipedia.org/wiki/Continuation-passing_style

Convert from

(+ (f 0 (g 1)) 2)

to

(g' (lambda (r0) (f' (lambda (r1) (+ r1 2)) 0 r0)) 1)

where data structure internally in Haskell is like

data AST = Node [AST] | Leaf Value
data Value = IntVal Int | Plus | Atom String | Lambda [String]

Implementation and description

import qualified Control.Monad.State as S

data AST = Node [AST] | Leaf Value
instance Show AST where
  show (Node xs) = "(" ++ unwords (map show xs) ++ ")"
  show (Leaf v) = show v

data Value = IntVal Int | Plus | Atom String | Lambda [String]
instance Show Value where
  show (IntVal i) = show i
  show Plus = "+"
  show (Atom name) = name
  show (Lambda names) = "lambda (" ++ unwords names ++ ")"

-- (+ (f 0 (g 1)) 2)
-- (g' (lambda (r0) (f' (lambda (r1) (+ r1 2)) 0 r0)) 1)
program :: AST
program = Node [Leaf Plus,
  Node [Leaf (Atom "f"), Leaf (IntVal 0), Node [Leaf (Atom "g"), Leaf (IntVal 1)]],
  Leaf (IntVal 2)]

main = do
  print program
  print $ cps program

cps :: AST -> AST
cps ast =
  let (newAst, modifiers) = S.runState (cps' ast) [] in
      foldl (flip ($)) newAst modifiers

cps' :: AST -> S.State [AST -> AST] AST
cps' (Node (Leaf (Atom f) : xs)) = do
  xs' <- mapM cps' xs
  n <- length `fmap` S.get
  let name = 'r' : show n
  append $ \root -> Node $
    (Leaf . Atom $ f ++ "'") :
    Node [Leaf (Lambda [name]), root] :
    xs'
  return $ Leaf (Atom name)
cps' (Node xs) = Node `fmap` mapM cps' xs
cps' c@(Leaf _) = return c

append x = S.modify (x :)

This converts correctly.

I used State Monad to modify given tree. The function cps starts state and the actual function cps' traverses given AST subtrees recursively.

(+ (f 0 (g 1)) 2)
   ^^^^^^^^^^^

When cps' sees this subtree, oh yes the first item of the list is a user-defined function and it's not tail-call, so cps' wants to replace the part with a new variable (say r), and enclose whole tree with new function f' and the arguments.

(f' (lambda (r) ((+ r 2) 0 (g 1))))
^^^^^^^^^^^^^^^^^   ^   ^^^^^^^^^^^

It's easy to change subtree but it's not trivial to change outside the subtree. But fortunately we already know that we only have to enclose something around the whole tree, so you can just save a function in state.

After cps' process is done, you apply all functions that the state has accumulatively to enclose trees. That will complete the job.