CPSLisp.hs

import qualified Control.Monad.State as S
import qualified Data.Map as M
import Control.Monad.Error ()

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 | Define String AST
  | Plus | Atom String | Lambda [String] AST
  | Halt
  | Begin [AST]
--  deriving Show
instance Show Value where
  show (IntVal i) = show i
  show (Define name body) = "define " ++ name ++ ' ': show body
  show Plus = "+"
  show Halt = "halt"
  show (Atom name) = name
  show (Lambda names x) =
    "(lambda (" ++ unwords names ++ ")" ++ " " ++ show x ++ ")"
  show (Begin xs) = "begin" ++ concatMap (("\n  " ++) . show) xs

type Env = M.Map String Value

-- (define f (lambda (x y) (+ x y)))
-- (define f' (lambda (k x y) (k (+ x y))))
--program1 :: AST
--program1 = Node [Leaf $ Define "f" (Leaf (Lambda ["x", "y"] (Node [Leaf Plus, Leaf (Atom "x"), Leaf (Atom "y")])))]

-- (define f (lambda (x) (+ x (f 1))))
--program1 = Node [Leaf $ Define "f" (Leaf (Lambda ["x"] (Node [Leaf Plus, Leaf (Atom "x"), Node [Leaf (Atom "f"), Leaf (IntVal 1)]])))]

-- (define f (lambda (x) (f x)))
--program1 = Node [Leaf $ Define "f" (Leaf (Lambda ["x"] (Node [Leaf (Atom "f"), Leaf (Atom "x")])))]

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

-- (+ (f 1 2) 3)
-- (f' (lambda (r0) (+ r0 3)) 1 2)
--program2 = Node [Leaf Plus, Node [Leaf (Atom "f"), Leaf (IntVal 1), Leaf (IntVal 2)], Leaf (IntVal 3)]

-- (f 1 2)
--program2 = Node [Leaf Plus, Node [Leaf (Atom "f"), Leaf (IntVal 0)], Leaf (IntVal 3)]

-- (f 2)
--program2 = Node [Leaf (Atom "f"), Leaf (IntVal 2)]

main :: IO ()
main = do
  -- (define call/cc (lambda (k f) (f k (lambda (d-k res) (k res)))))
  let callcc = Node [
        Leaf (Define "call/cc'"
          (Leaf (Lambda ["k", "f"] (Node [
            Leaf (Atom "f"),
            Leaf (Atom "k"),
            Leaf (Lambda ["d-k", "res"] (Node [
              Leaf (Atom "k"),
              Leaf (Atom "res")]))]))))]
  print callcc
  -- (+ 10 (call/cc (lambda (c) (+ 2 (c 1)))))
  let caller = Node [
        Leaf Plus,
        Leaf (IntVal 10),
        Node [
          Leaf (Atom "call/cc"),
          Leaf (Lambda ["c"] (Node [
            Leaf Plus,
            Leaf (IntVal 2),
            Node [
              Leaf (Atom "c"),
              Leaf (IntVal 1)
            ] ])) ] ]
  print caller
  print $ encloseHalt caller
  let program = Node [Leaf $ Begin [callcc, cps $ encloseHalt caller]]
  print program
  print $ eval program
  --print program1
  --print $ encloseHalt program2
  --putStrLn ""
  ----print $ eval $ Node [Leaf $ Begin [program1, program2]]
  --let program = Node [Leaf $ Begin [cps program1, cps $ encloseHalt program2]]
  --print program
  --putStrLn ""
  --print $ eval program
  print "very cool."

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

encloseHalt :: AST -> AST
encloseHalt root = Node [Leaf Halt, root]

cps0 :: AST -> S.State [AST -> AST] AST
cps0 x@(Leaf _) = return x
cps0 (Node (Leaf (Define name (Leaf (Lambda args body))) : [])) = do
  let name' = name ++ "'"
  let body' = cps body -- not cps'
  return $ Node [Leaf $ Define name' $ Leaf (Lambda ("k" : args) (Node [Leaf (Atom "k"), body']))]
cps0 (Node xs) = Node `fmap` mapM cps' xs

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 ++ "'") :
    Leaf (Lambda [name] root) :
    xs'
  return $ Leaf (Atom name)
-- cps' (Node [Leaf (Begin xs)]) =
--   (Node . (: []) . Leaf . Begin) `fmap` mapM cps' xs
cps' (Node xs) = Node `fmap` mapM cps' xs
cps' c@(Leaf _) = return c

append :: (AST -> AST) -> S.State [AST -> AST] ()
append x = S.modify (x :)

type RuntimeState a = S.StateT Env (Either (a, Env)) a

eval :: AST -> Value
eval ast = fst $ either id id $ S.runStateT (eval' ast) M.empty

eval' :: AST -> RuntimeState Value
eval' (Node [Leaf (Define name body)]) = do
  body' <- eval' body
  S.modify $ M.insert name body'
  return $ Atom name
eval' (Node [Leaf (Begin xs)]) = do
  xs' <- mapM eval' xs
  return $ last xs'
eval' (Node (Leaf Plus : xs)) = do
  xs' <- mapM eval' xs
  return $ foldl1 f xs'
  where
    f (IntVal a) (IntVal b) = IntVal (a + b)
    f _ _ = error "argument error: + Int Int"
eval' (Node (f : xs)) = do
  f' <- eval' f
  xs' <- mapM eval' xs
  val <- apply f' xs'
  state <- S.get
  S.lift $ Left (val, state)
eval' (Leaf (Atom x)) = do
  x' <- M.lookup x `fmap` S.get
  maybe (error . ("cannot find " ++) . show $ x) return x'
eval' (Leaf i) = return i
eval' x = error $ show x

apply :: Value -> [Value] -> RuntimeState Value
apply (Lambda names body) args = do
  mapM_ S.modify $ zipWith M.insert names args
  eval' body
--apply Halt (arg:_) = error $ show arg -- w
apply Halt (arg:_) = return arg
apply f _ = error $ show f ++ " is not lambda!"