LlvmMangler: Make sure no symbols slip through re-.typing
[ghc.git] / compiler / llvmGen / LlvmMangler.hs
1 -- -----------------------------------------------------------------------------
2 -- | GHC LLVM Mangler
3 --
4 -- This script processes the assembly produced by LLVM, rearranging the code
5 -- so that an info table appears before its corresponding function.
6 --
7
8 module LlvmMangler ( llvmFixupAsm ) where
9
10 import DynFlags ( DynFlags )
11 import ErrUtils ( showPass )
12 import LlvmCodeGen.Ppr ( infoSection )
13
14 import Control.Exception
15 import Control.Monad ( when )
16 import qualified Data.ByteString.Char8 as B
17 import Data.Char
18 import System.IO
19
20 import Data.List ( sortBy )
21 import Data.Function ( on )
22
23 #if x86_64_TARGET_ARCH
24 #define REWRITE_AVX
25 #endif
26
27 -- Magic Strings
28 secStmt, infoSec, newLine, textStmt, dataStmt, syntaxUnified :: B.ByteString
29 secStmt = B.pack "\t.section\t"
30 infoSec = B.pack infoSection
31 newLine = B.pack "\n"
32 textStmt = B.pack "\t.text"
33 dataStmt = B.pack "\t.data"
34 syntaxUnified = B.pack "\t.syntax unified"
35
36 infoLen :: Int
37 infoLen = B.length infoSec
38
39 -- Search Predicates
40 isType :: B.ByteString -> Bool
41 isType = B.isPrefixOf (B.pack "\t.type")
42
43 -- section of a file in the form of (header line, contents)
44 type Section = (B.ByteString, B.ByteString)
45
46 -- | Read in assembly file and process
47 llvmFixupAsm :: DynFlags -> FilePath -> FilePath -> IO ()
48 llvmFixupAsm dflags f1 f2 = {-# SCC "llvm_mangler" #-} do
49 showPass dflags "LLVM Mangler"
50 r <- openBinaryFile f1 ReadMode
51 w <- openBinaryFile f2 WriteMode
52 ss <- readSections r w
53 hClose r
54 let fixed = (map rewriteAVX . fixTables) ss
55 mapM_ (writeSection w) fixed
56 hClose w
57 return ()
58
59 rewriteSymType :: B.ByteString -> B.ByteString
60 rewriteSymType s =
61 foldl (\s' (typeFunc,typeObj)->replace typeFunc typeObj s') s types
62 where
63 types = [ (B.pack "@function", B.pack "@object")
64 , (B.pack "%function", B.pack "%object")
65 ]
66
67 -- | Splits the file contents into its sections
68 readSections :: Handle -> Handle -> IO [Section]
69 readSections r w = go B.empty [] []
70 where
71 go hdr ss ls = do
72 e_l <- (try (B.hGetLine r))::IO (Either IOError B.ByteString)
73
74 -- Note that ".type" directives at the end of a section refer to
75 -- the first directive of the *next* section, therefore we take
76 -- it over to that section.
77 let (tys, ls') = span isType ls
78 cts = rewriteSymType $ B.intercalate newLine $ reverse ls'
79
80 -- Decide whether to directly output the section or append it
81 -- to the list for resorting.
82 let finishSection
83 | infoSec `B.isInfixOf` hdr =
84 cts `seq` return $ (hdr, cts):ss
85 | otherwise =
86 writeSection w (hdr, cts) >> return ss
87
88 case e_l of
89 Right l | l == syntaxUnified
90 -> finishSection >>= \ss' -> writeSection w (l, B.empty)
91 >> go B.empty ss' tys
92 | any (`B.isPrefixOf` l) [secStmt, textStmt, dataStmt]
93 -> finishSection >>= \ss' -> go l ss' tys
94 | otherwise
95 -> go hdr ss (l:ls)
96 Left _ -> finishSection >>= \ss' -> return (reverse ss')
97
98 -- | Writes sections back
99 writeSection :: Handle -> Section -> IO ()
100 writeSection w (hdr, cts) = do
101 when (not $ B.null hdr) $
102 B.hPutStrLn w hdr
103 B.hPutStrLn w cts
104
105 #if REWRITE_AVX
106 rewriteAVX :: Section -> Section
107 rewriteAVX = rewriteVmovaps . rewriteVmovdqa
108
109 rewriteVmovdqa :: Section -> Section
110 rewriteVmovdqa = rewriteInstructions vmovdqa vmovdqu
111 where
112 vmovdqa, vmovdqu :: B.ByteString
113 vmovdqa = B.pack "vmovdqa"
114 vmovdqu = B.pack "vmovdqu"
115
116 rewriteVmovap :: Section -> Section
117 rewriteVmovap = rewriteInstructions vmovap vmovup
118 where
119 vmovap, vmovup :: B.ByteString
120 vmovap = B.pack "vmovap"
121 vmovup = B.pack "vmovup"
122 #else /* !REWRITE_AVX */
123 rewriteAVX :: Section -> Section
124 rewriteAVX = id
125 #endif /* !REWRITE_SSE */
126
127 rewriteInstructions :: B.ByteString -> B.ByteString -> Section -> Section
128 rewriteInstructions matchBS replaceBS (hdr, cts) =
129 (hdr, replace matchBS replaceBS cts)
130
131 replace :: B.ByteString -> B.ByteString -> B.ByteString -> B.ByteString
132 replace matchBS replaceBS = loop
133 where
134 loop :: B.ByteString -> B.ByteString
135 loop cts =
136 case B.breakSubstring matchBS cts of
137 (hd,tl) | B.null tl -> hd
138 | otherwise -> hd `B.append` replaceBS `B.append`
139 loop (B.drop (B.length matchBS) tl)
140
141 -- | Reorder and convert sections so info tables end up next to the
142 -- code. Also does stack fixups.
143 fixTables :: [Section] -> [Section]
144 fixTables ss = map strip sorted
145 where
146 -- Resort sections: We only assign a non-zero number to all
147 -- sections having the "STRIP ME" marker. As sortBy is stable,
148 -- this will cause all these sections to be appended to the end of
149 -- the file in the order given by the indexes.
150 extractIx hdr
151 | B.null a = 0
152 | otherwise = 1 + readInt (B.takeWhile isDigit $ B.drop infoLen a)
153 where (_,a) = B.breakSubstring infoSec hdr
154
155 indexed = zip (map (extractIx . fst) ss) ss
156
157 sorted = map snd $ sortBy (compare `on` fst) indexed
158
159 -- Turn all the "STRIP ME" sections into normal text sections, as
160 -- they are in the right place now.
161 strip (hdr, cts)
162 | infoSec `B.isInfixOf` hdr = (textStmt, cts)
163 | otherwise = (hdr, cts)
164
165 -- | Read an int or error
166 readInt :: B.ByteString -> Int
167 readInt str | B.all isDigit str = (read . B.unpack) str
168 | otherwise = error $ "LLvmMangler Cannot read " ++ show str
169 ++ " as it's not an Int"
170