Better compiler interface

3 files changed, 227 insertions, 48 deletions

diff --git a/Main.hs b/Main.hs
index eb84072..dc76f65 100644
--- a/Main.hs
+++ b/Main.hs
@@ -1,14 +1,19 @@
 module Main where
 
+import Control.Exception (bracket)
+import Control.Monad
+import System.Directory (removeFile)
+import System.Environment
 import System.Exit
 import System.IO
 import System.Process
-import Debug.Trace
+import qualified System.Posix.Temp as Posix
 
 import BuildIR
 import CodeGen
 import Defs
 import Optimiser
+import OptionParser
 import Pretty
 import ProgramParser
 import TypeCheck
@@ -22,35 +27,103 @@ f <?> pre = \a -> case f a of
     Right x -> Right x
 
 
-tracePrettyId :: Pretty a => a -> a
-tracePrettyId x = trace (pretty x) x
-
 eitherToIO :: Either String a -> IO a
 eitherToIO = either die return
 
-performCompile :: String -> IO ()
-performCompile source = do
-    let eres = return source
-                 >>= parseProgram <?> "Parse error"
-                 -- >>= return . tracePrettyId
-                 >>= typeCheck <?> "Type error"
-                 >>= buildIR <?> "IR building error"
-                 >>= optimise <?> "Error while optimising"
-                 -- >>= return . traceShowId
-                 >>= verify <?> "Verify error"
-                 >>= return . tracePrettyId
-                 >>= codegen <?> "Codegen error"
-
-    asm <- eitherToIO eres
-    -- hPutStr stderr asm
+extensionForStage :: CompilerStage -> String
+extensionForStage StageTypeCheck = ".lang"
+extensionForStage StageIR = ".ir"
+extensionForStage StageAsm = ".asm"
+extensionForStage StageObject = ".o"
+extensionForStage StageExecutable = ""
 
-    writeFile "z_output.asm" asm
+inputFile :: String -> IO String
+inputFile "-" = getContents
+inputFile fname = readFile fname
 
-    hPutStrLn stderr "Assembling with yasm..."
-    callCommand "yasm -w+all -fmacho64 z_output.asm -o z_output.o"
+outputFile :: String -> String -> IO ()
+outputFile "-" str = putStr str
+outputFile fname str = writeFile fname str
 
-    hPutStrLn stderr "Linking with ld..."
-    callCommand "ld z_output.o liblang.o -o z_output"
+mkTempFile :: IO FilePath
+mkTempFile = do
+    (path, handle) <- Posix.mkstemp "/tmp/tmp.lang."
+    hClose handle
+    return path
 
 main :: IO ()
-main = getContents >>= performCompile
+main = do
+    opts <- getArgs >>= eitherToIO . optionParser
+
+    when (oShowHelp opts) $ do
+        name <- getProgName
+        putStrLn $ "Usage: " ++ name ++ " [options] files..."
+        putStrLn ""
+        putStrLn "-h, --help         Show this help"
+        putStrLn "-o, --output FILE  Send output to the specified file instead of z_output[.ext]."
+        putStrLn "-O, --optimise     Optimise the output. (Happens at IR stage.)"
+        putStrLn "--stage typecheck  Only parse and typecheck the program, outputting a pretty-"
+        putStrLn "                   printed version."
+        putStrLn "--stage ir         Compile until (possibly optimised) IR, then output that."
+        putStrLn "-S, --stage asm    Compile until assembly, then output that."
+        putStrLn "-c, --stage obj    Compile to an object file."
+        putStrLn "--stage exe        Default mode; compile and link."
+        exitSuccess
+
+    source <- case oSourceFiles opts of
+        [] -> die "No source files given. Run with --help for usage information."
+        [name] -> inputFile name
+        _ -> die "Only one source file supported at the moment."
+
+    when (oOutputName opts == "-" && oLastStage opts >= StageObject) $
+        die "Will not write binary file to stdout."
+
+    let outname = if oAppendExtension opts
+                    then oOutputName opts ++ extensionForStage (oLastStage opts)
+                    else oOutputName opts
+
+    let laststage = oLastStage opts
+
+    ast' <- eitherToIO $ (parseProgram <?> "Parse error") source
+    ast <- eitherToIO $ (typeCheck <?> "Type error") ast'
+    when (laststage == StageTypeCheck) $ do
+        outputFile outname $ pretty ast
+        exitSuccess
+
+    ir <- eitherToIO $ (buildIR <?> "IR building error") ast
+    iropt' <- eitherToIO $ (optimise (oOptimise opts) <?> "Error while optimising") ir
+    iropt <- eitherToIO $ (verify <?> "Verify error") iropt'
+    when (laststage == StageIR) $ do
+        outputFile outname $ pretty iropt
+        exitSuccess
+
+    asm <- eitherToIO $ (codegen <?> "Codegen error") iropt
+    when (laststage == StageAsm) $ do
+        outputFile outname asm
+        exitSuccess
+
+
+    let objfnameIO = case laststage of
+            StageObject -> return outname
+            StageExecutable -> mkTempFile
+            _ -> undefined
+
+    let rmObjfileIO = case outname of
+            "-" -> const $ return ()
+            _ -> removeFile
+
+    bracket objfnameIO rmObjfileIO $ \objfname -> do
+        hPutStrLn stderr "Assembling with yasm..."
+        let yasmprocspec = (proc "yasm" ["-w+all", "-fmacho64", "-", "-o", objfname]) {std_in = CreatePipe}
+        yasmcode <- withCreateProcess yasmprocspec $ \(Just pipe) _ _ ph -> do
+            hPutStr pipe asm
+            hFlush pipe
+            hClose pipe
+            waitForProcess ph
+        case yasmcode of
+            ExitSuccess -> return ()
+            ExitFailure _ -> die "yasm failed!"
+        when (laststage == StageObject) exitSuccess
+
+        hPutStrLn stderr "Linking with ld..."
+        callProcess "ld" [objfname, "liblang.o", "-o", outname]
diff --git a/Optimiser.hs b/Optimiser.hs
index f94441c..c438c82 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -1,4 +1,4 @@
-module Optimiser(optimise) where
+module Optimiser(optimise, OptimiserLevel(..)) where
 
 import Data.Either
 import Data.Function
@@ -14,35 +14,33 @@ import ReplaceRefs
 import Utils
 
 
+data OptimiserLevel = Level0 | Level1
+  deriving (Show, Eq, Ord)
+
 type Optimisation = IRProgram -> IRProgram
 type FuncOptimisation = IRFunc -> IRFunc
 
-optimise :: IRProgram -> Error IRProgram
-optimise prog =
-    let optlist = [trace "-- OPT PASS --" {-, \p -> trace (pretty p) p-}] ++ optimisations
+optimise :: OptimiserLevel -> IRProgram -> Error IRProgram
+optimise optlevel prog =
+    let optlist = {-[trace "-- OPT PASS --", \p -> trace (pretty p) p] ++ -} optimisations
         reslist = scanl (flip ($)) prog $ cycle optlist
         passreslist = map fst $ filter (\(_, i) -> i `mod` length optlist == 0) $ zip reslist [0..]
         applyFinalOpts p = foldl (flip ($)) p finaloptimisations
-    in if True
-        then return $ applyFinalOpts $ 
-                fst $ fromJust $ find (uncurry (==)) $ zip passreslist (tail passreslist)
-        else return $ reslist !! 5
+    in return $ applyFinalOpts $ 
+            fst $ fromJust $ find (uncurry (==)) $ zip passreslist (tail passreslist)
   where
-    -- optimisations = map funcopt
-    --     [chainJumps, removeUnusedBlocks]
-    -- optimisations = map funcopt
-    --     [chainJumps, mergeTerminators, looseJumps,
-    --      removeUnusedBlocks,
-    --      constantPropagate, movPush]
-    optimisations = map funcopt
-        [chainJumps, mergeTerminators, looseJumps,
-         removeUnusedBlocks, removeDuplicateBlocks,
-         identityOps,
-         constantPropagate, movPush,
-         arithPush, removeUnusedInstructions,
-         evaluateInstructions, evaluateTerminators]
-    finaloptimisations = map funcopt
-        [reorderBlocks, flipJccs, invertJccs]
+    optimisations = case optlevel of
+        Level0 -> map funcopt [chainJumps, mergeTerminators, looseJumps, removeUnusedBlocks]
+        Level1 -> map funcopt
+                    [chainJumps, mergeTerminators, looseJumps,
+                     removeUnusedBlocks, removeDuplicateBlocks,
+                     identityOps,
+                     constantPropagate, movPush,
+                     arithPush, removeUnusedInstructions,
+                     evaluateInstructions, evaluateTerminators]
+    finaloptimisations = case optlevel of
+        Level0 -> map funcopt [flipJccs]
+        Level1 -> map funcopt [reorderBlocks, flipJccs, invertJccs]
 
 
 funcopt :: FuncOptimisation -> Optimisation
diff --git a/OptionParser.hs b/OptionParser.hs
new file mode 100644
index 0000000..ef0a813
--- /dev/null
+++ b/OptionParser.hs
@@ -0,0 +1,108 @@
+{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}
+module OptionParser(optionParser, Options(..), CompilerStage(..)) where
+
+import Control.Monad
+import Control.Monad.Except
+import Control.Monad.State.Strict
+import Data.List
+import Data.Maybe
+import qualified Data.Map.Strict as Map
+import Debug.Trace
+
+import Defs (Error)
+import Optimiser (OptimiserLevel(..))
+
+
+data Options = Options
+    { oOutputName :: String
+    , oSourceFiles :: [String]
+    , oLastStage :: CompilerStage
+    , oShowHelp :: Bool
+    , oAppendExtension :: Bool
+    , oOptimise :: OptimiserLevel
+    }
+  deriving Show
+
+initOptions :: Options
+initOptions = Options
+    { oOutputName = "z_output"
+    , oSourceFiles = []
+    , oLastStage = StageExecutable
+    , oShowHelp = False
+    , oAppendExtension = True
+    , oOptimise = Level0
+    }
+
+data CompilerStage = StageTypeCheck | StageIR | StageAsm | StageObject | StageExecutable
+  deriving (Show, Eq, Ord, Enum)
+
+
+data DBEntry
+    = WithArgument (String -> Options -> Error Options)
+    | NoArgument (Options -> Error Options)
+    | MaybeArgument (Maybe String -> Options -> Error Options)
+
+optionDatabase :: Map.Map String DBEntry
+optionDatabase = Map.fromList $ map swap . concatMap sequence . map swap $
+    [ (["--help", "-h"], NoArgument $ \o -> return $ o {oShowHelp = True})
+    , (["--output", "-o"], WithArgument $ \s o -> return $ o {oOutputName = s, oAppendExtension = False})
+    , (["-c"], NoArgument $ \o -> return $ o {oLastStage = StageObject})
+    , (["-S"], NoArgument $ \o -> return $ o {oLastStage = StageAsm})
+    , (["--optimise", "-O"], MaybeArgument $ \ms o -> case ms of
+            Nothing -> return $ o {oOptimise = Level1}
+            Just "0" -> return $ o {oOptimise = Level0}
+            Just "1" -> return $ o {oOptimise = Level1}
+            Just val -> Left $ "Invalid optimiser level '" ++ val ++ "'")
+    , (["--stage"], WithArgument $ \s o -> case s of
+            "typecheck" -> return $ o {oLastStage = StageTypeCheck}
+            "ir" -> return $ o {oLastStage = StageIR}
+            "asm" -> return $ o {oLastStage = StageAsm}
+            "obj" -> return $ o {oLastStage = StageObject}
+            "exe" -> return $ o {oLastStage = StageExecutable}
+            _ -> Left $ "Unrecognised value for '--stage'")
+    ]
+
+swap :: (a, b) -> (b, a)
+swap (a, b) = (b, a)
+
+
+optionParser :: [String] -> Error Options
+optionParser args = foldM foldfunc initOptions kvs
+  where
+    kvs = collectKeyValues args
+    foldfunc opts (arg@('-':_), mval) = case (Map.lookup arg optionDatabase, mval) of
+        (Nothing, _) -> Left $ "Unknown command-line option '" ++ arg ++ "'"
+        (Just (NoArgument f), Nothing) -> f opts
+        (Just (NoArgument _), Just val) ->
+            Left $ "Option '" ++ arg ++ "' takes no argument, but got '" ++ val ++ "'"
+        (Just (WithArgument f), Just val) -> f val opts
+        (Just (WithArgument _), Nothing) ->
+            Left $ "Option '" ++ arg ++ "' takes an argument, but got none"
+        (Just (MaybeArgument f), _) -> f mval opts
+    foldfunc opts (arg, Nothing) = return $ opts {oSourceFiles = oSourceFiles opts ++ [arg]}
+    foldfunc _ (_, Just _) = undefined
+
+collectKeyValues :: [String] -> [(String, Maybe String)]
+collectKeyValues [] = []
+collectKeyValues (arg@('-':'-':_:_):rest) = case findIndex (== '=') arg of
+    Nothing -> case Map.lookup arg optionDatabase of
+        Nothing -> (arg, Nothing) : collectKeyValues rest
+        Just (NoArgument _) -> (arg, Nothing) : collectKeyValues rest
+        Just (WithArgument _) -> case rest of
+            [] -> [(arg, Nothing)]
+            (hd:tl) -> (arg, Just hd) : collectKeyValues tl
+        Just (MaybeArgument _) -> (arg, Nothing) : collectKeyValues rest
+    Just idx -> (take idx arg, Just $ drop (idx + 1) arg) : collectKeyValues rest
+collectKeyValues (arg@('-':_:_:_):rest) = case Map.lookup (take 2 arg) optionDatabase of
+    Nothing -> (take 2 arg, Just $ drop 2 arg) : collectKeyValues rest
+    Just (NoArgument _) -> (take 2 arg, Nothing) : collectKeyValues (('-' : drop 2 arg) : rest)
+    Just (WithArgument _) -> (take 2 arg, Just $ drop 2 arg) : collectKeyValues rest
+    Just (MaybeArgument _) -> (take 2 arg, Just $ drop 2 arg) : collectKeyValues rest
+collectKeyValues (arg@['-',_]:rest) = case Map.lookup arg optionDatabase of
+    Nothing -> (arg, Nothing) : collectKeyValues rest
+    Just (NoArgument _) -> (arg, Nothing) : collectKeyValues rest
+    Just (WithArgument _) -> case rest of
+        [] -> [(arg, Nothing)]
+        (next:rest2) -> (arg, Just next) : collectKeyValues rest2
+    Just (MaybeArgument _) -> (arg, Nothing) : collectKeyValues rest
+collectKeyValues (arg:rest) = (arg, Nothing) : collectKeyValues rest