Defensify naked read in LLVM mangler

[ghc-hetmet.git] / compiler / llvmGen / LlvmMangler.hs

diff --git a/compiler/llvmGen/LlvmMangler.hs b/compiler/llvmGen/LlvmMangler.hs
index 390f4a8..27d2a84 100644 (file)
--- a/compiler/llvmGen/LlvmMangler.hs
+++ b/compiler/llvmGen/LlvmMangler.hs
@@ -19,12 +19,16 @@ import qualified Data.ByteString.Char8 as BS
 
 import LlvmCodeGen.Ppr ( infoSection, iTableSuf )
 
+import Data.Char
+import Outputable
+import Util
+
 
 {- Configuration. -}
 newSection, oldSection, functionSuf, tableSuf, funDivider, eol :: ByteString
 newSection  = BS.pack "\n.text\n"
 oldSection  = BS.pack infoSection
-functionSuf = BS.pack "_info:"
+functionSuf = BS.pack $ if ghciTablesNextToCode then "_info:" else "\n_"
 tableSuf    = BS.pack $ "_info" ++ iTableSuf ++ ":"
 funDivider  = BS.pack "\n\n"
 eol         = BS.pack "\n"
@@ -97,18 +101,18 @@ oneTable f str =
     let last' xs = if (null xs) then 0 else last xs
 
         -- get the function
-        (bl, al) = BS.breakSubstring functionSuf str
-        start = last' $ BS.findSubstrings funDivider bl
+        (bl, al)          = BS.breakSubstring functionSuf str
+        start             = last' $ BS.findSubstrings funDivider bl
         (before, fheader) = BS.splitAt start bl
-        (fun, after) = BS.breakSubstring funDivider al
-        label = snd $ BS.breakEnd eolPred bl
+        (fun, after)      = BS.breakSubstring funDivider al
+        label             = snd $ BS.breakEnd eolPred bl
 
         -- get the info table
-        ilabel = label `BS.append` tableSuf
-        (bit, itable) = BS.breakSubstring ilabel after
-        (itable', ait) = BS.breakSubstring funDivider itable
-        istart = last' $ BS.findSubstrings funDivider bit
-        (bit', iheader) = BS.splitAt istart bit
+        ilabel            = label `BS.append` tableSuf
+        (bit, itable)     = BS.breakSubstring ilabel after
+        (itable', ait)    = BS.breakSubstring funDivider itable
+        istart            = last' $ BS.findSubstrings funDivider bit
+        (bit', iheader)   = BS.splitAt istart bit
 
         -- fixup stack alignment
         fun' = fixupStack fun BS.empty
@@ -124,21 +128,26 @@ oneTable f str =
               BS.appendFile f bl
               return BS.empty
 
-          else if BS.null itable
-                  then error $ "Function without matching info table! ("
-                              ++ (BS.unpack label) ++ ")"
+          else if ghciTablesNextToCode
+                  then if BS.null itable
+                          then error $ "Function without matching info table! ("
+                                      ++ (BS.unpack label) ++ ")"
+                          else do
+                              mapM_ (BS.appendFile f) function
+                              return remainder
 
                   else do
-                      mapM_ (BS.appendFile f) function
-                      return remainder
+                      -- TNTC not turned on so just fix up stack
+                      mapM_ (BS.appendFile f) [before, fheader, fun']
+                      return after
 
 -- | Replace the current section in a function or table header with the
 -- text section specifier.
 replaceSection :: ByteString -> ByteString
 replaceSection sec =
     let (s1, s2) = BS.breakSubstring oldSection sec
-        s1' = fst $ BS.breakEnd eolPred s1
-        s2' = snd $ BS.break eolPred s2
+        s1'      = fst $ BS.breakEnd eolPred s1
+        s2'      = snd $ BS.break eolPred s2
     in s1' `BS.append` newSection `BS.append` s2'
 
 
@@ -155,10 +164,11 @@ replaceSection sec =
 fixupStack :: ByteString -> ByteString -> ByteString
 fixupStack fun nfun | BS.null nfun =
     let -- fixup sub op
-        (a, b) = BS.breakSubstring (BS.pack ", %esp\n") fun
-        (a', num) = BS.breakEnd dollarPred a
-        num' = BS.pack $ show (read (BS.unpack num) + 4::Int)
-        fix = a' `BS.append` num'
+        (a, b)       = BS.breakSubstring (BS.pack ", %esp\n") fun
+        (a', strNum) = BS.breakEnd dollarPred a
+        Just num     = readInt (BS.unpack strNum)
+        num'         = BS.pack $ show (num + 4::Int)
+        fix          = a' `BS.append` num'
     in if BS.null b
           then nfun `BS.append` a
           else fixupStack b (nfun `BS.append` fix)
@@ -168,15 +178,28 @@ fixupStack fun nfun =
         (a, b) = BS.breakSubstring (BS.pack "jmp") fun
         -- We need to avoid processing jumps to labels, they are of the form:
         -- jmp\tL..., jmp\t_f..., jmpl\t_f..., jmpl\t*%eax...
-        labelJump = BS.index b 4 == 'L'
-        (jmp, b') = BS.break eolPred b
-        (a', numx) = BS.breakEnd dollarPred a
-        (num, x) = BS.break commaPred numx
-        num' = BS.pack $ show (read (BS.unpack num) + 4::Int)
-        fix = a' `BS.append` num' `BS.append` x `BS.append` jmp
+        labelJump   = BS.index b 4 == 'L'
+        (jmp, b')   = BS.break eolPred b
+        (a', numx)  = BS.breakEnd dollarPred a
+        (strNum, x) = BS.break commaPred numx
+        Just num    = readInt (BS.unpack strNum)
+        num'        = BS.pack $ show (num + 4::Int)
+        fix         = a' `BS.append` num' `BS.append` x `BS.append` jmp
     in if BS.null b
           then nfun `BS.append` a
           else if labelJump
                 then fixupStack b' (nfun `BS.append` a `BS.append` jmp)
                 else fixupStack b' (nfun `BS.append` fix)
 
+
+-- | 'read' is one of my least favourite functions.
+readInt :: String -> Maybe Int
+readInt str
+       | not $ null $ filter (not . isDigit) str
+       = pprTrace "LLvmMangler"
+               (text "Cannot read" <+> text (show str) <+> text "as it's not an Int")
+               Nothing
+
+       | otherwise
+       = Just $ read str
+