1 files changed, 38 insertions, 21 deletions

diff --git a/Main.hs b/Main.hs
index dc76f65..3e2a325 100644
--- a/Main.hs
+++ b/Main.hs
@@ -2,16 +2,19 @@ module Main where
 
 import Control.Exception (bracket)
 import Control.Monad
+import qualified Data.ByteString as BS
 import System.Directory (removeFile)
 import System.Environment
 import System.Exit
 import System.IO
 import System.Process
 import qualified System.Posix.Temp as Posix
+import qualified System.Info as System
 
 import BuildIR
 import CodeGen
 import Defs
+import LibLang
 import Optimiser
 import OptionParser
 import Pretty
@@ -51,6 +54,9 @@ mkTempFile = do
     hClose handle
     return path
 
+withTempFile :: (FilePath -> IO a) -> IO a
+withTempFile = bracket mkTempFile removeFile
+
 main :: IO ()
 main = do
     opts <- getArgs >>= eitherToIO . optionParser
@@ -103,27 +109,38 @@ main = do
         exitSuccess
 
 
-    let objfnameIO = case laststage of
-            StageObject -> return outname
-            StageExecutable -> mkTempFile
+    let withObjfile = case laststage of
+            StageObject -> ($ outname)
+            StageExecutable -> withTempFile
             _ -> undefined
 
-    let rmObjfileIO = case outname of
-            "-" -> const $ return ()
-            _ -> removeFile
-
-    bracket objfnameIO rmObjfileIO $ \objfname -> do
+    let yasmFormat = case System.os of
+            "linux" -> "elf64"
+            "darwin" -> "macho64"
+            os -> error $ "Your OS (" ++ os ++ ") is unknown, can't create binary"
+
+    let yasmToFile outfname writer = do
+            let yasmprocspec =
+                    (proc "yasm" ["-w+all", "-f" ++ yasmFormat, "-", "-o", outfname])
+                        {std_in = CreatePipe}
+            yasmcode <- withCreateProcess yasmprocspec $ \(Just pipe) _ _ ph -> do
+                _ <- writer pipe
+                hFlush pipe
+                hClose pipe
+                waitForProcess ph
+            case yasmcode of
+                ExitSuccess -> return ()
+                ExitFailure _ -> die "yasm failed!"
+
+    withObjfile $ \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]
+        yasmToFile objfname (\pipe -> hPutStr pipe asm)
+
+        case laststage of
+            StageObject -> return ()
+            StageExecutable ->
+                withTempFile $ \liblangfname -> do
+                    yasmToFile liblangfname (\pipe -> BS.hPut pipe libLangSource)
+                    hPutStrLn stderr "Linking with ld..."
+                    callProcess "ld" [objfname, liblangfname, "-o", outname]
+            _ -> undefined