Lambda Lifting as SLPJ - 1987 Book Page 227

LambdaLifting.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DeriveFunctor #-}

module LambdaLifting where

import Control.Monad.Trans.Reader
import Control.Monad.Trans.Writer
import Control.Monad.State hiding (fix)
import Data.Char
import Data.List

fix :: ((a -> b) -> a -> b) -> a -> b
fix f = f (fix f)

data Fix f = In { out :: f (Fix f) }

instance Show (f (Fix f)) => Show (Fix f) where
  show (In f) = "{" ++ show f ++ "}"

cata :: Functor f => (f a -> a) -> (Fix f -> a) -> Fix f -> a
cata psi f = psi . fmap f . out

cataRec :: Functor f => (f a -> a) -> Fix f -> a
cataRec psi = fix (cata psi)

data ExpF a = Var String
            | App a [a]
            | Lam [String] a
            | Lit Int
            | Op String a a
            | Let String a a
            deriving (Functor, Show)

type Exp = Fix ExpF

app :: Exp -> [Exp] -> Exp
app e1 e2 = In (App e1 e2)

var :: String -> Exp
var x = In (Var x)

lam :: [String] -> Exp -> Exp
lam vs b = In (Lam vs b)

lit :: Int -> Exp
lit n = In (Lit n)

op :: String -> Exp -> Exp -> Exp
op s e1 e2 = In (Op s e1 e2)

leT :: String -> Exp -> Exp -> Exp
leT s e1 e2 = In (Let s e1 e2)

without :: Eq a => [a] -> [a] -> [a]
without = foldr (filter . (/=)) -- Like \\ but removes all occurrences

freeVars :: Exp -> [String]
freeVars = cataRec alg
   where alg (Var v)    = [v]
         alg (App l r)  = l ++ concat r
         alg (Lam vs e) = e `without` vs
         alg (Lit _)    = []
         alg (Op _ e1 e2) = e1 ++ e2

applyTo :: Exp -> [String] -> Exp
applyTo = foldl (\e a -> app e $ [var a])

transformLam :: [String] -> Exp -> Exp
transformLam globals = cataRec alg
   where alg (Lam vs e) =
           let vars = freeVars e `without` (globals ++ vs)
            in
              if null vars then lam vs e else lam vars (lam vs e) `applyTo` vars
         alg e = In e

data Def = Def String [String] Exp
               deriving Show

type ClosM = WriterT [Def] (ReaderT [String] (State Int))

localCtx :: ([String] -> [String]) -> ClosM a -> ClosM a
localCtx f m = WriterT(local f (runWriterT m))

askCtx :: ClosM [String]
askCtx = lift ask


gen :: ClosM String
gen = do x <- get
         modify (+ 1)
         return ("$" ++ [chr (ord 'F' + x)])

liftLam :: Exp -> ClosM Exp
liftLam = cataRec alg
   where alg (Var v)      = return $ var v
         alg (App l r)    = do x <- l
                               y <- sequence r
                               return $ app x y
         alg (Op s e1 e2) = do x <- e1
                               y <- e2
                               return $ op s x y
         alg (Lit i)      = return $ lit i
         alg (Lam vs e)  = do
           ctx <- askCtx
           v <- e
           case ctx of
             ["let"] -> return $ lam vs v
             _ -> do fresh <- gen
                     _ <- tell (return (Def fresh vs v))
                     return $ var fresh

         alg (Let v e1 e2) = do
            x <- localCtx (\x -> x ++ ["let"]) e1
            y <- e2
            return $ leT v x y

uncurryLambdasAndApps :: Exp -> Exp
uncurryLambdasAndApps = cataRec alg
 where alg (Lam vs (In (Lam ws e))) = lam (vs ++ ws) e
       alg (App (In (App e1 e2)) e3) = app e1 (e2 ++ e3)
       alg e = In e

lambdaLifting :: [String] -> Exp -> (Exp, [Def])
lambdaLifting globals e = evalState (runReaderT (runWriterT (liftLam(uncurryLambdasAndApps(transformLam globals e)))) []) 0

toCommaSep :: [String] -> String
toCommaSep [x] = x
toCommaSep xs = concat ["(", intercalate "," xs, ")"]

pretty :: Exp -> String
pretty = cataRec alg
  where alg (Var v) = v
        alg (App l r) = concat ["(", l, ") (", toCommaSep r,  ")"]
        alg (Lam vs e) = concat ["\\", toCommaSep vs, ".(", e, ")"]
        alg (Op n x y) = concat ["(", x, " ", n, " ", y, ")"]
        alg (Lit n) = show n
        alg (Let v e1 e2) = concat ["let ", v, " = ", e1, " in ", e2]

sample :: Exp
sample = app (lam ["y"] (app (var "f") [lam ["x"] (var "y")])) [lit 5]

after :: Exp
after = transformLam ["f"] sample

lifted :: (Exp, [Def])
lifted = lambdaLifting ["f"] sample

main :: IO()
main  = do putStrLn $ pretty sample
           putStrLn $ pretty after
           mapM_ (\(Def n vs e) ->
             putStrLn(concat [n,  " ", toCommaSep vs, " = ", pretty e])) (snd lifted)
           putStrLn $ pretty (fst lifted)