Frank.hs

{-# LANGUAGE GADTs #-}
module Main where

import Parser
import Compile
import TypeCheck
import Syntax
import RefineSyntax
import DesugarSyntax
import Debug

import qualified Shonky.Semantics as Shonky

import Control.Monad

import Data.IORef
import qualified Data.Map.Strict as M
import Data.List

import System.Console.CmdArgs.Explicit
import System.Directory
import System.Environment
import System.Exit
import System.IO

type Args = [(String,[String])]

type PMap = M.Map FilePath (Prog Raw)

-- Function checks whether or not a main function exists in the program
existsMain :: Prog t -> Bool
existsMain (MkProg xs) = any (== "main") [id | DefTm (Def id _ _ _) _ <- xs]

splice :: Prog Raw -> Tm Raw -> Prog Raw
splice (MkProg xs) tm = MkProg $ xs ++ ys
  where ys = [sig, cls]
        sig = SigTm (Sig id (CType [] peg b) b) b
        peg = Peg ab ty b
        ty = TVar "%X" b
        ab = Ab (AbVar "£" b) (ItfMap M.empty (Raw Generated)) b
        cls = ClsTm (MHCls id (Cls [] tm b) b) b
        id = "%eval"
        b = Raw Generated

--TODO: LC: Fix locations?
exorcise :: Prog Desugared -> (Prog Desugared, TopTm Desugared)
exorcise (MkProg xs) = (prog, DefTm (head evalDef) a)
  where prog = MkProg (map (swap DataTm a) dts ++
                       map (swap ItfTm a) itfs ++
                       map (swap DefTm a) hdrs)
        dts = [d | DataTm d _ <- xs]
        itfs = [i | ItfTm i _ <- xs]
        defs = [d | DefTm d _ <- xs]
        (evalDef,hdrs) = partition isEvalTm defs
        isEvalTm (Def id _ _ _) = id == "%eval"
        a = Desugared Generated

extractEvalUse :: TopTm Desugared -> Use Desugared
extractEvalUse (DefTm (Def _ _ [cls] _) _) = getBody cls
  where getBody (Cls [] (Use u _) _) = u
        getBody _ = error "extractEvalUse: eval body invariant broken"

glue :: Prog Desugared -> TopTm Desugared -> Prog Desugared
glue (MkProg xs) x = MkProg (x : xs)

parseProg :: FilePath -> Args -> IO (Either String (Prog Raw))
parseProg fileName args =
  do m <- pProg (Right M.empty) fileName
     case m of
       Left msg  -> return $ Left msg
       Right map -> return $ Right $ MkProg $ M.foldl combine [] map
  where pProg :: Either String PMap -> FilePath -> IO (Either String PMap)
        pProg (Left msg) _ = return $ Left msg
        pProg (Right map) fname | M.member fname map = return $ Right map
        pProg (Right map) fname =
          let ext = if isSuffixOf ".fk" fname then "" else ".fk" in
          do m <- parseFile (fname ++ ext) incPaths
             case m of
               Left msg -> return $ Left msg
               Right (p,fs) -> foldM pProg (Right (M.insert fname p map)) fs

        combine :: [TopTm Raw] -> Prog Raw -> [TopTm Raw]
        combine xs (MkProg ys) = xs ++ ys

        incPaths :: [String]
        incPaths = case lookup "inc" args of
          Just xs -> xs
          Nothing -> []

        parseFile :: String -> [FilePath] ->
                     IO (Either String (Prog Raw,[String]))
        parseFile name incs = let fs = name : map (++ name) incs in
          do m <- foldM g Nothing fs
             case m of
               Nothing ->
                 die ("could not find " ++ name ++
                      " in search path:" ++ (concat $ intersperse "," fs))
               Just x -> runTokenProgParse <$> readFile x
          where g :: Maybe FilePath -> FilePath -> IO (Maybe FilePath)
                g Nothing x = do b <- doesFileExist x
                                 return $ if b then Just x else Nothing
                g x@(Just _) _ = return x

parseEvalTm :: String -> IO (Tm Raw)
parseEvalTm v = case runTokenParse tm v of
  Left err -> die err
  Right tm -> return tm

refineComb :: Either String (Prog Raw) -> Tm Raw -> IO (Prog Refined)
refineComb prog tm = case prog of
    Left err -> die err
    Right p -> case refine (splice p tm) of
      Left err -> die err
      Right p' -> return p'

refineAndDesugarProg :: Either String (Prog Raw) -> IO (Prog Desugared)
refineAndDesugarProg p =
  case p of
    Left err -> die err
    Right p -> case refine p of
      Left err -> die err
      Right p' -> return $ desugar p'

checkProg :: Prog Desugared -> Args -> IO (Prog Desugared)
checkProg p _ =
  case check p of
    Left err -> die err
    Right p' -> return p'

checkUse :: Prog Desugared -> Use Desugared -> IO (Use Desugared, VType Desugared)
checkUse p use =
  case inferEvalUse p use of
    Left err -> die err
    Right (use, ty) -> return (use, ty)

compileProg :: String -> Prog Desugared -> Args -> IO Shonky.Env
compileProg progName p args =
  if ("output-shonky",[]) `elem` args then
    do compileToFile p progName
       Shonky.loadFile progName
  else return $ Shonky.load $ compile p

evalProg :: Shonky.Env -> String -> IO ()
evalProg env tm =
  case Shonky.try env tm of
    Shonky.Ret v -> putStrLn $ (show . Shonky.ppVal) v
    comp -> do -- putStrLn $ "Generated computation: " ++ show comp
               v <- Shonky.ioHandler comp
               putStrLn $ (show . Shonky.ppVal) v

compileAndRunProg :: String -> Args -> IO ()
compileAndRunProg fileName args =
  do let progName = takeWhile (/= '.') fileName
     prog <- parseProg fileName args
     case lookup "eval" args of
       Just [v] -> do tm <- parseEvalTm v
                      -- lift tm to top term and append to prog
                      p <- refineComb prog tm
                      -- tear apart again
                      let (p',ttm) = exorcise (desugar p)
                          use = extractEvalUse ttm
                      p'' <- checkProg p' args
                      (use', _) <- checkUse p'' use
                      env <- compileProg progName (glue p'' ttm) args
                      evalProg env "%eval()"
       Just _ -> die "only one evaluation point permitted"
       Nothing -> do p <- refineAndDesugarProg prog
                     p' <- checkProg p args
                     env <- compileProg progName p' args
                     case lookup "entry-point" args of
                       Just [v] -> evalProg env (v ++ "()")
                       Just _  -> die "only one entry point permitted"
                       Nothing ->
                         if existsMain p' then
                           evalProg env "main()"
                         else
                           putStrLn ("Compilation successful! " ++
                                     "(no main function defined)")
arguments :: Mode Args
arguments =
  mode "frank" [] "Frank program" (flagArg (upd "file") "FILE")
  [flagNone ["output-shonky"] (("output-shonky",[]):)
   "Output Shonky code"
  ,flagNone ["debug-output"] (("debug-output",[]):)
   "Enable all debugging facilities"
  ,flagNone ["debug-verbose"] (("debug-verbose",[]):)
   "Enable verbose variable names etc. on output"
  ,flagNone ["debug-parser"] (("debug-parser",[]):)
    "Enable output of parser logs"
  ,flagNone ["debug-tc"] (("debug-tc",[]):)
   "Enable output of type-checking logs"
  ,flagReq ["include", "I"] updInc "INC"
   "Add a search path for source files"
  ,flagReq ["eval"] (upd "eval") "USE"
   "Run Frank computation USE (default: main!)"
  ,flagReq ["entry-point"] (upd "entry-point") "NAME"
   "Run computation NAME (default: main)"
  ,flagHelpSimple (("help",[]):)]
  where upd msg x v = Right $ (msg,[x]):v
        updInc x args = case lookup "inc" args of
          Just xs -> Right $ ("inc",x:xs):args
          Nothing -> Right $ ("inc",[x]):args

-- handy for testing in ghci
run f = compileAndRunProg f []

main :: IO ()
main = do
  hSetBuffering stdin NoBuffering
  hSetEcho stdin False
  args <- processArgs arguments
  let debugVerboseOn =   ("debug-output",[]) `elem` args ||
                         ("debug-verbose", []) `elem` args
      debugParserOn =    ("debug-output",[]) `elem` args ||
                         ("debug-parser", []) `elem` args
      debugTypeCheckOn = ("debug-output",[]) `elem` args ||
                         ("debug-tc", []) `elem` args
  writeIORef debugMode (debugVerboseOn, debugParserOn, debugTypeCheckOn)
  if ("help",[]) `elem` args then
     print $ helpText [] HelpFormatDefault arguments
  else case lookup "file" args of
    Just [f] -> compileAndRunProg f args
    Just _  -> die "too many Frank sources specified."
    Nothing -> die "no Frank source specified."