module ContCompiler where 

import Language.Haskell.Pretty

import Target
import ContType

comp = compK

compK :: Expr -> Target
compK (Const c)         = TUnitM c
compK (Var x)           = TUnitM (TVar x)
compK (Val (x, m) n)    = compK m `TBindM` TLambda1 x 
                         (compK n)
compK (Let decls n)     = TLet (map (\ (x, m) -> let TUnitM c = compK m in (PVar x, c)) decls)
                               (compK n)
compK (App f x)         = compK f `TBindM` TLambda1 "_f"
                         (compK x `TBindM` TLambda1 "_x"
                         (TApp1 (TVar "_f") (TVar "_x")))
compK (Lambda x m)      = TUnitM (if x=="_" then TLambda0 (compK m) else TLambda1 x (compK m)) 
compK (Delay m)         = TUnitM (compK m)
compK (If e1 e2 e3)     = compK e1 `TBindM` 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          = TApp1 (TVar "abortK") (TApp1 (TVar "_while") (TVar "_break"))
compK (Goto l)          = mkGoto l

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

-- compCC: uses call/cc for while/goto.
compCC :: Expr -> Target
compCC (Const c)         = TUnitM c
compCC (Var x)           = TUnitM (TVar x)
compCC (Val (x, m) n)    = compCC m `TBindM` TLambda1 x 
                        (compCC n)
compCC (Let decls n)     = TLet (map (\ (x, m) -> let TUnitM c = compCC m in (PVar x, c)) decls)
                              (compCC n)
compCC (App f x)         = compCC f `TBindM` TLambda1 "_f"
                        (compCC x `TBindM` TLambda1 "_x"
                        (TApp1 (TVar "_f") (TVar "_x")))
compCC (Lambda x m)      = TUnitM (if x=="_" then TLambda0 (compCC m) else TLambda1 x (compCC m)) 
compCC (Delay m)         = TUnitM (compCC m)
compCC (If e1 e2 e3)     = compCC e1 `TBindM` 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          = TApp1 (TVar "abortK") (TApp1 (TVar "_while") (TVar "_break"))
compCC (Goto l)          = TVar l

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


mkGoto l  = TApp1 (TVar "abortK") (TApp1 (TVar l) (TVar "()"))

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

compCCWhile e1 e2 = TApp1 (TVar "callccK") (TLambda1 "_break"
                   (TLet [(PVar "_while", body)]
                         (TVar "_while")))
  where body = compCC e1 `TBindM` TLambda1 "_b" 
               (TIf (TVar "_b") (compCC e2 `TBindM` TLambda0 (TVar "_while"))
                                (TUnitM (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 `TBindM` (TLambda0
                                      (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  `TBindM` (TLambda0
                                     (TApp1 (TVar "callccK") (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 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 `TBindM` TLambda0 endl)]
    aux ((l1,t1):(l2,t2):lts) = (PVar l1, t1 `TBindM` TLambda0 (TVar l2)) : aux ((l2,t2):lts)


compKExprs []      = TUnitM (TVar "()")
compKExprs [e]     = compK e
compKExprs (e:es)  = compK e `TBindM` TLambda0 (compKExprs es)

compCCExprs []     = TUnitM (TVar "()")
compCCExprs [e]    = compCC e
compCCExprs (e:es) = compCC e `TBindM` TLambda0 (compCCExprs es)