module ContMain where

import Language.Haskell.Pretty

import Target
import ContType
import ContParser
import ContCompiler

primitives = map (\ n -> (n, mkOp n)) ["**", "*", "+", "-", "++", ":", "==", "/=", "<", "<=", ">=", ">"]
          ++ map (\ n -> (n, divOp n)) ["/", "%", "//"] 
          ++ map (\ n -> (n, mkPrimN 2 n)) ["mplus","set", "throw"]
          ++ map (\ n -> (n, mkPrimN 1 n)) 
               ["get", "writeStr", "write"]
          ++ map (\ n -> (n, mkPrimN 0 n)) ["readChar", "eof"]
          ++ map (\ n -> (n, mkPureN 1 n)) ["Trick", "unTrick", "Mu", "unMu",
               "not", "null", "head", "tail", "show", "atoi", "atof"]

simplifyK :: Target -> Target
simplifyK (m `TBind` k) = simplifyK (TLambda1 "c" (TApp1 (shift 1 0 m)
                                         (TLambda1 "a" (TApp1 (TApp1 (shift 2 0 k) (TBound 0))
                                                              (TBound 1)))))
simplifyK (TUnit t)      = simplifyK (TLambda1 "c" (TApp1 (TBound 0) (shift 1 0 t)))
simplifyK (TLambda1 v t) = TLambda1 v (simplifyK t)
simplifyK (TLambda0 t)   = TLambda0 (simplifyK t)
simplifyK (TLambda p t)  = TLambda (map simplifyKPat p) (simplifyK t)
simplifyK (TApp (TVar "KST") [x])    = simplifyK x
simplifyK (TApp (TVar "unKST") [x])    = simplifyK x
simplifyK (TApp (TVar "KST") (x:xs)) = TApp (simplifyK x) (map simplifyK xs)
simplifyK (TApp (TVar "unKST") (x:xs)) = TApp (simplifyK x) (map simplifyK xs)
simplifyK (TApp f xs)    = TApp (simplifyK f) (map simplifyK xs)
simplifyK (TApp1 (TVar "KST") x)    = simplifyK x
simplifyK (TApp1 (TVar "unKST") x)    = simplifyK x
simplifyK (TApp1 f x)    = TApp1 (simplifyK f) (simplifyK x)
simplifyK (TInfix l o r) = TInfix (simplifyK l) o (simplifyK r)
simplifyK (TLet ds t)    = TLet (map (\ (p,t) -> (simplifyKPat p,simplifyK t)) ds) (simplifyK t)
simplifyK (TIf c t e)    = TIf (simplifyK c) (simplifyK t) (simplifyK e)
simplifyK (TCase t alts) = TCase (simplifyK t) (map (\ (p,t) -> (simplifyKPat p,simplifyK t)) alts)
simplifyK t              = t

simplifyKPat (PCon "KST" [p]) = simplifyKPat p
simplifyKPat (PCon v ps)      = PCon v (map simplifyKPat ps)
simplifyKPat p                = p

simplifyKDecls :: [(t, Target)] -> [(t, Target)]
simplifyKDecls ds = map (\ (p,t) -> (p,simplifyK t)) ds

test = testK
testDecls = testKDecls

testK useDo prog = prettyPrint $ translateToHs useDo
           	 $ removeSugar False $ simplify $ instOnce primitives
           	 $ deBruijn $ compK $ myParse prog

testKDecls useDo prog = concat $ map (\ str -> str ++ "\n") $ map prettyPrint $ translateToHsDecls useDo
           	      $ removeSugarDecls False $ simplifyDecls $ instOnceDecls primitives
           	      $ deBruijnDecls $ compKDecls $ myParseDecls prog

testCC useDo prog = prettyPrint $ translateToHs useDo
            	  $ removeSugar False $ simplify $ instOnce primitives
            	  $ deBruijn $ compCC $ myParse prog

testCCDecls useDo prog = concat $ map (\ str -> str ++ "\n") $ map prettyPrint $ translateToHsDecls useDo
           	       $ removeSugarDecls False $ simplifyDecls $ instOnceDecls primitives
           	       $ deBruijnDecls $ compCCDecls $ myParseDecls prog

envCC = [("abort",  TApp1 (TVar "KST") (TLambda1 "v" (TLambda0 (TBound 0)))),
         ("callcc", TApp1 (TVar "KST") (TLambda1 "h"
                          (TLambda [PVar "c"]
                              (TApp (TBound 1)
                                    [TLambda1 "a" (TLambda0 (TApp1 (TBound 1) (TBound 0))),
                                                                          TBound 0]))))]


toCPS = toCPSK
toCPSDecls = toCPSKDecls

toCPSK prog = prettyPrint $ translateToHs False $ removeSugar True
            $ simplify $ simplifyK $ instOnce envCC 
            $ simplify $ instOnce primitives $ deBruijn $ compK $ myParse prog

toCPSKDecls prog = concat $ map (\ str -> str ++ "\n") $ map prettyPrint $ translateToHsDecls False $ removeSugarDecls True
                 $ simplifyDecls $ instOnceDecls envCC $ simplifyKDecls
                 $ simplifyDecls $ instOnceDecls primitives $ deBruijnDecls
                 $ compKDecls $ myParseDecls prog


toCPSCC prog = prettyPrint $ translateToHs False $ removeSugar True
             $ simplify $ simplifyK $ instOnce envCC 
             $ simplify $ instOnce primitives $ deBruijn $ compCC $ myParse prog

toCPSCCDecls prog = concat $ map (\ str -> str ++ "\n") $ map prettyPrint $ translateToHsDecls False $ removeSugarDecls True
                   $ simplifyDecls $ instOnceDecls envCC $ simplifyKDecls
                   $ simplifyDecls $ instOnceDecls primitives $ deBruijnDecls
                   $ compCCDecls $ myParseDecls prog

utilGoto = 
  unlines [ "begin"
          , "  set xP 1;"
          , "label1:"
          , "  if get xP > 100 then"
          , "    goto label2"
          , "  else ();"
          , "  set xP (get xP * 2);"
          , "  goto label1;"
          , "label2:"
          , "  get xP"
          , "end"
          ]

utilWhileDecl = 
  unlines [ "foo = \\ n ->"
          , "  begin"
          , "    set xP 1;"
          , "    set yP n;"
          , "    while get yP > 0 do"
          , "      begin"
          , "        val r = get xP in"
          , "        val n = get yP in"
          , "        begin"
          , "          if n==10 then break"
          , "          else if n==3 then"
          , "            begin"
          , "              set yP (n-1);"
          , "              continue"
          , "            end"
          , "          else ();"
          , "          set xP (r*n);"
          , "          set yP (n-1)"
          , "        end"
          , "      end;"
          , "    get xP"
          , "  end"
          ]
utilWhile = "let " ++ utilWhileDecl ++ "in foo 9\n"

utilFactDecl = "fact = \\ n -> if n==0 then 1 else n*fact(n-1)"
utilFact = "let " ++ utilFactDecl ++ " in fact 9"

utilBazDecl = 
  unlines [ "baz = \\ x ->"
          , "  callcc (\\ k ->"
          , "    100 + (if x==0 then 1 else k x))"
          ]
utilBaz = "let " ++ utilBazDecl ++ " in baz 1"

utilMultDecl =
  unlines [ "multlist = \\ list ->" 
          , "  let aux = \\ xs -> \\ k -> begin"
          , "        set xP 1; set yP xs;"
          , "        while not (null (get yP)) do begin"
          , "          val y = get yP in"
          , "          val n = head y in"
          , "          if n == 0 then k 0 else"
          , "            begin set xP (get xP*n); set yP (tail y); writeStr \" \"; write n end"
          , "        end;"
          , "        get xP"
          , "    end in"
          , "  val result = callcc (\\ k -> aux list k)  in begin"
          , "    writeStr \"; result = \";"
          , "    write result"
          , "  end"
          ]

utilMult = "let " ++ utilMultDecl ++ " in multlist (1:2:3:0:4:5:[])"

-- for example,
-- putStrLn $ test False utilGoto
-- putStrLn $ test False utilWhile
-- putStrLn $ test False utilBaz
-- putStrLn $ test False utilMult
-- putStrLn $ toCPS utilFact