module ContCompiler where 

import Language.Haskell.Pretty

import Target
import ContType

comp = compK

compK :: Expr -> Target
compK (Const c)         = TReturn c
compK (Var x)           = TReturn (TVar x)
compK (Val (x, m) n)    = compK m `TBind` TLambda1 x 
                         (compK n)
compK (Let decls n)     = TLet (map (\ (x, m) -> let TReturn c = compK m in (PVar x, c)) decls)
                               (compK n)
compK (App f x)         = compK f `TBind` TLambda1 "_f"
                         (compK x `TBind` TLambda1 "_x"
                         (TApp1 (TVar "_f") (TVar "_x")))
compK (Lambda x m)      = TReturn (if x=="_" then TLambda0 (compK m) else TLambda1 x (compK m)) 
compK (Delay m)         = TReturn (compK m)
compK (If e1 e2 e3)     = compK e1 `TBind` TLambda1 "_b"
                         (TIf (TVar "_b") (compK e2) (compK e3))
compK (While e1 e2)     = compKWhile e1 e2
compK (Begin es)        = compKBegin es 
compK Break             = mkGoto "_break" "()"
compK Continue          = mkGoto "_while" "_break"
compK (Goto l)          = mkGoto l "()"

compKDecls :: Decls -> [(Pattern, Target)]
compKDecls decls = map (\ (x, m) -> let TReturn c = compK m in (PVar x, c)) decls

-- compCC: uses call/cc for while/goto.
compCC :: Expr -> Target
compCC (Const c)         = TReturn c
compCC (Var x)           = TReturn (TVar x)
compCC (Val (x, m) n)    = compCC m `TBind` TLambda1 x 
                        (compCC n)
compCC (Let decls n)     = TLet (map (\ (x, m) -> let TReturn c = compCC m in (PVar x, c)) decls)
                              (compCC n)
compCC (App f x)         = compCC f `TBind` TLambda1 "_f"
                        (compCC x `TBind` TLambda1 "_x"
                        (TApp1 (TVar "_f") (TVar "_x")))
compCC (Lambda x m)      = TReturn (if x=="_" then TLambda0 (compCC m) else TLambda1 x (compCC m)) 
compCC (Delay m)         = TReturn (compCC m)
compCC (If e1 e2 e3)     = compCC e1 `TBind` TLambda1 "_b"
                        (TIf (TVar "_b") (compCC e2) (compCC e3))
compCC (While e1 e2)     = compCCWhile e1 e2
compCC (Begin es)        = compCCBegin es 
compCC Break             = mkGoto "_break" "()"
compCC Continue          = mkGoto "_while" "_break"
compCC (Goto l)          = TVar l

compCCDecls :: Decls -> [(Pattern, Target)]
compCCDecls decls = map (\ (x, m) -> let TReturn c = compCC m in (PVar x, c)) decls


mkGoto l v  = TApp1 (TVar "abort") (TApp1 (TVar l) (TVar v))

compKWhile e1 e2 = TApp1 (TVar "KST")
                     (TLambda1 "_break"
                       (TLet [(PCon "KST" [PVar "_while"], body)]
                             (TApp1 (TVar "_while") (TVar "_break"))))
  where body = compK e1 `TBind` TLambda1 "_b" 
               (TIf (TVar "_b") (compK e2 `TBind` TLambda0 (TApp1 (TVar "KST") (TVar "_while")))
                                (TReturn (TVar "()")))

compCCWhile e1 e2 = TApp1 (TVar "callcc") (TLambda1 "_break"
                   (TLet [(PVar "_while", body)]
                         (TVar "_while")))
  where body = compCC e1 `TBind` TLambda1 "_b" 
               (TIf (TVar "_b") (compCC e2 `TBind` TLambda0 (TVar "_while"))
                                (TReturn (TVar "()")))

splitLabel []                 = ([], [])
splitLabel ((Nothing, e):les) = let (es,ls) = splitLabel les in (e:es, ls)
splitLabel ((Just l, e):les)  = let (es,ls) = splitLabel les in ([], (l, e:es):ls)

compKBegin les = let (es, less) = splitLabel les
                   in case less of 
                        []         -> compKExprs es
                        ((l1,_):_) -> compKExprs es `TBind` (TLambda0
                                      (TApp1 (TVar "KST") (TLambda1 "_end"
                                        (TLet (compKLabeledExprs less "_end")
                                              (TApp1 (TVar l1) (TVar "()"))))))

compCCBegin les = let (es, less) = splitLabel les
                    in case less of   
                       []         -> compCCExprs es
                       ((l1,_):_) -> compCCExprs es  `TBind` (TLambda0
                                     (TApp1 (TVar "callcc") (TLambda1 "_end"
                                       (TLet (compCCLabeledExprs less (TApp1 (TVar "_end") (TVar "()")))
                                             (TVar l1)))))

compKLabeledExprs les endl = let tes = map (\(l,es) -> (l, compKExprs es)) les
                               in aux tes
  where
    addLabel t l = TLambda0 (TApp1 (TApp1 (TVar "unKST") t) (TVar l))
    aux [(l,t)]               = [(PVar l, addLabel t endl)]
    aux ((l1,t1):(l2,t2):lts) = (PVar l1, addLabel t1 l2) : aux ((l2,t2):lts)

compCCLabeledExprs les endl = let tes = map (\(l,es) -> (l, compCCExprs es)) les
                                in aux tes
  where
    aux [(l,t)]               = [(PVar l, t `TBind` TLambda0 endl)]
    aux ((l1,t1):(l2,t2):lts) = (PVar l1, t1 `TBind` TLambda0 (TVar l2)) : aux ((l2,t2):lts)


compKExprs []      = TReturn (TVar "()")
compKExprs [e]     = compK e
compKExprs (e:es)  = compK e `TBind` TLambda0 (compKExprs es)

compCCExprs []     = TReturn (TVar "()")
compCCExprs [e]    = compCC e
compCCExprs (e:es) = compCC e `TBind` TLambda0 (compCCExprs es)