1 files changed, 98 insertions, 25 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]