Let generated code work with yasm on linux

2 files changed, 31 insertions, 9 deletions

diff --git a/Compiler.hs b/Compiler.hs
index eab7058..f02a50a 100644
--- a/Compiler.hs
+++ b/Compiler.hs
@@ -5,6 +5,7 @@ module Compiler(compile) where
 import Control.Monad.State.Strict
 import Control.Monad.Writer.Strict
 import Data.List
+import qualified System.Info as System (os)
 
 import AST
 
@@ -30,6 +31,19 @@ genId :: LineWriter Int
 genId = state $ \i -> (i, i + 1)
 
 
+mangleName :: String -> String
+mangleName f = case System.os of
+    "darwin" -> '_' : f
+    "linux" -> f
+    _ -> undefined
+
+generateCall :: String -> String
+generateCall f = case System.os of
+    "darwin" -> "call " ++ mangleName f
+    "linux" -> "call [" ++ mangleName f ++ " wrt ..got]"
+    _ -> undefined
+
+
 compile :: Program -> String
 compile (Program inss) = runLineWriter $ do
     emit0 prologue
@@ -67,12 +81,12 @@ compileIns (ILoop inss off) = do
     emit $ "jnz .Lbd_" ++ show loopid
     emit0 $ ".Laf_" ++ show loopid ++ ":"
 compileIns (IInput off) = do
-    emit "call _getchar"
+    emit $ generateCall "getchar"
     emit $ "mov [" ++ cursorOffset off ++ "], al"
 compileIns (IOutput off) = do
     emit "xor edi, edi"
     emit $ "mov dil, [" ++ cursorOffset off ++ "]"
-    emit "call _putchar"
+    emit $ generateCall "putchar"
 compileIns IStart = return ()
 
 isTwoPower :: Byte -> Maybe Int
@@ -81,19 +95,20 @@ isTwoPower v = findIndex (==v) (take 8 $ iterate (* 2) 1)
 
 prologue :: String
 prologue =
-         "global _main"
-    ++ "\nextern _calloc, _free, _putchar, _getchar, _write"
+         "global " ++ mangleName "main"
+    ++ "\nextern " ++ intercalate ", "
+            (map mangleName ["calloc", "free", "putchar", "getchar", "write"])
     ++ "\ndefault rel"
     ++ "\nsection .text"
 
-    ++ "\n_main:"
+    ++ "\n" ++ mangleName "main" ++ ":"
     ++ "\n\tpush rbp"
     ++ "\n\tmov rbp, rsp"
     ++ "\n\tpush rbx"
     ++ "\n\tsub rsp, 8"
     ++ "\n\tmov edi, " ++ show tapeLength
     ++ "\n\tmov esi, 1"
-    ++ "\n\tcall _calloc"
+    ++ "\n\t" ++ generateCall "calloc"
     ++ "\n\ttest rax, rax"
     ++ "\n\tjz .allocation_failure"
     ++ "\n\tlea rbx, [rax + " ++ show (tapeLength `div` 2) ++ "]  ; rbx = cursor"
@@ -103,7 +118,7 @@ prologue =
 epilogue :: String
 epilogue =
        "\n\tmov rdi, [rsp]"
-    ++ "\n\tcall _free"
+    ++ "\n\t" ++ generateCall "free"
     ++ "\n\txor eax, eax"
     ++ "\n.return:"
     ++ "\n\tadd rsp, 8"
@@ -122,7 +137,7 @@ epilogue =
     ++ "\n\tmov edi, 2"
     ++ "\n\tmov rsi, rsp"
     ++ "\n\tmov edx, 20"
-    ++ "\n\tcall _write"
+    ++ "\n\t" ++ generateCall "write"
     ++ "\n\tadd rsp, 32"
     ++ "\n\tmov eax, 1"
     ++ "\n\tjmp .return"
diff --git a/Main.hs b/Main.hs
index 9c33f83..c7ca464 100644
--- a/Main.hs
+++ b/Main.hs
@@ -4,6 +4,7 @@ import Control.Monad
 import Data.Char
 import System.Environment
 import System.Exit
+import qualified System.Info as System (os)
 import System.Process
 
 import AST
@@ -40,5 +41,11 @@ main = do
             interpret opt (map (fromIntegral . ord) input) >>= (putStr . map (chr . fromIntegral))
         EMCompile -> do
             writeFile (fname ++ ".asm") $ compile opt
-            callProcess "yasm" ["-f", "macho64", fname ++ ".asm", "-o", fname ++ ".o"]
+
+            let format = case System.os of
+                            "darwin" -> "macho64"
+                            "linux" -> "elf64"
+                            _ -> undefined
+
+            callProcess "yasm" ["-f", format, fname ++ ".asm", "-o", fname ++ ".o"]
             callProcess "gcc" [fname ++ ".o", "-o", fname ++ ".exe"]