Some CPS transformations.

CPS.hs

-- see: http://matt.might.net/articles/cps-conversion

{-# OPTIONS_GHC -Wall #-}
{-# 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

-- Data.Unique.Extended -------------------------------------------------------

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

-- Lambda ---------------------------------------------------------------------

-- expression; lambda expression; etc.

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

instance PP.Pretty Expr where
  pretty expr = case expr of
    Lam v e ->
             PP.text "(λ " <> PP.textStrict v <> PP.text ". "
      PP.<$> PP.indent 2 (PP.pretty e) <> PP.text ")"

    Var v -> PP.textStrict v

    App e f -> PP.text "(" <> PP.pretty e PP.<+> PP.pretty f <> ")"

data AExpr =
    AVar Text
  | ALam [Text] CExpr

instance PP.Pretty AExpr where
  pretty aexpr = case aexpr of
    AVar v -> PP.textStrict v

    ALam vs cexpr ->
      let pvs = PP.text "(" <> PP.hcat (fmap PP.pretty vs) <> PP.text ")"
      in  PP.text "(λ " <> pvs <> PP.text ". "
          PP.<$> PP.indent 2 (PP.pretty cexpr) <> PP.text ")"

data CExpr =
    CApp AExpr [AExpr]

instance PP.Pretty CExpr where
  pretty cexpr = case cexpr of
    CApp aexpr aexprs ->
      let paexprs = fmap PP.pretty aexprs
      in  PP.text "(" <> PP.pretty aexpr
            PP.<+> PP.fillSep paexprs <> PP.text ")"

-- naive CPS xform ------------------------------------------------------------

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"

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

-- higher-order CPS xform -----------------------------------------------------

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

-- hybrid CPS xform -----------------------------------------------------------

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

-- test -----------------------------------------------------------------------

test0 :: Expr
test0 = App (Var "g") (Var "a")

res0 :: IO CExpr
res0 = nt test0 (AVar "halt")

res1 :: IO CExpr
res1 = hot test0 (\ans -> return (CApp (AVar "halt") [ans]))

res2 :: IO CExpr
res2 = tc test0 (AVar "halt")

main :: IO ()
main = do
  print (PP.pretty test0)

  putStrLn mempty

  r0 <- res0
  print (PP.pretty r0)

  putStrLn mempty

  r1 <- res1
  print (PP.pretty r1)

  putStrLn mempty

  r2 <- res2
  print (PP.pretty r2)