module ContCompiler where

import Language.Haskell.Pretty

import Target
import Expr

comp = compK

compK :: Expr -> Target
compK (Const c)         = TReturn c
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 (LetGlobal ds n)  =
  let len = length ds
      rs  = mkRefs len
      (ms, ps) = unzip $ zipWith (\ (x, m) r -> (m, (x, r))) ds rs
      cs  = map (\ m -> let TReturn c = compK m in c) ms
      s0  = TApp (TVar $ mkAt $ mkTupleCon len) cs
    in TApp (TVar "resetGlobal")
          [s0, TLet (map (\ (x, r) -> (PVar x, TApp (TVar "cmpPos") [TVar "p2_1", r])) ps) (compK n)]
compK (LetLocal ds n)   =
  let len = length ds
      rs  = mkRefs len
      (ms, ps) = unzip $ zipWith (\ (x, m) r -> (m, (x, r))) ds rs
      cs  = map (\ m -> let TReturn c = compK m in c) ms
      s0  = TApp (TVar $ mkAt $ mkTupleCon len) cs
    in TApp (TVar "resetLocal")
          [s0, TLet (map (\ (x, r) -> (PVar x, TApp (TVar "cmpPos") [TVar "p2_2", r])) ps) (compK n)]
compK (Var x)
    | isMutVar x        = TApp1 (TVar "get") (TVar x)
    | otherwise         = TReturn (TVar x)
compK (App (App (Var ":=") (Var x)) m)
                        = compK m                      `TBind` TLambda1 "_y"
                          (TApp1 (TVar "set") (TVar x) `TBind` TLambda1 "_f"
                          (TApp1 (TVar "_f") (TVar "_y")))
compK (App (Var "&") (Var x))
    | isMutVar x        = TReturn (TVar x)
compK (App f x)         = compK f `TBind` TLambda1 "_f"
                         (compK x `TBind` TLambda1 "_x"
                         (TApp1 (TVar "_f") (TVar "_x")))                          
compK (Lambda "_" m)    = TReturn (TLambda0 (compK m))
compK (Lambda x m)      = TReturn (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 (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 (LetGlobal ds n)  =
  let len = length ds
      rs  = mkRefs len
      (ms, ps) = unzip $ zipWith (\ (x, m) r -> (m, (x, r))) ds rs
      cs  = map (\ m -> let TReturn c = compCC m in c) ms
      s0  = TApp (TVar $ mkAt $ mkTupleCon len) cs
    in TApp (TVar "resetGlobal")
          [s0, TLet (map (\ (x, r) -> (PVar x, TApp (TVar "cmpPos") [TVar "p2_1", r])) ps) (compCC n)]
compCC (LetLocal ds n)   =
  let len = length ds
      rs  = mkRefs len
      (ms, ps) = unzip $ zipWith (\ (x, m) r -> (m, (x, r))) ds rs
      cs  = map (\ m -> let TReturn c = compCC m in c) ms
      s0  = TApp (TVar $ mkAt $ mkTupleCon len) cs
    in TApp (TVar "resetLocal")
          [s0, TLet (map (\ (x, r) -> (PVar x, TApp (TVar "cmpPos") [TVar "p2_2", r])) ps) (compCC n)]
compCC (Var x)
    | isMutVar x         = TApp1 (TVar "get") (TVar x)
    | otherwise          = TReturn (TVar x)
compCC (App (App (Var ":=") (Var x)) m)
                         = compCC m                     `TBind` TLambda1 "_y"
                           (TApp1 (TVar "set") (TVar x) `TBind` TLambda1 "_f"
                           (TApp1 (TVar "_f") (TVar "_y")))
compCC (App (Var "&") (Var x))
    | isMutVar x         = TReturn (TVar x)
compCC (App f x)         = compCC f  `TBind` TLambda1 "_f"
                           (compCC x `TBind` TLambda1 "_x"
                           (TApp1 (TVar "_f") (TVar "_x")))
compCC (Lambda "_" m)    = TReturn (TLambda0 (compCC m))
compCC (Lambda x m)      = TReturn (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 "KIO")
                     (TLambda1 "_break"
                       (TLet [(PCon "KIO" [PVar "_while"], body)]
                             (TApp1 (TVar "_while") (TVar "_break"))))
  where body = compK e1 `TBind` TLambda1 "_b"
               (TIf (TVar "_b") (compK e2 `TBind` TLambda0 (TApp1 (TVar "KIO") (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,_):_) -> TApp1 (TVar "KIO") (TLambda1 "_end"
                                        (TLet (compKLabeledExprs less "_end")
                                              (TApp (TVar "unKIO") [compKExprs es, TVar l1])))

compCCBegin les = let (es, less) = splitLabel les
                    in case less of
                       []         -> compCCExprs es
                       ((l1,_):_) -> TApp1 (TVar "callcc") (TLambda1 "_end"
                                      (TLet (compCCLabeledExprs less (TApp1 (TVar "_end") (TVar "()")))
                                            (compCCExprs es  `TBind` (TLambda0 (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 "unKIO") t) (TVar l)) -- \ () -> unKIO t 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)