summaryrefslogtreecommitdiff
path: root/compiler/llvmGen/LlvmMangler.hs
blob: 8652a890cf83981c8b1641a4bd6e52e40ebcb68b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
{-# LANGUAGE CPP #-}

-- -----------------------------------------------------------------------------
-- | GHC LLVM Mangler
--
-- This script processes the assembly produced by LLVM, rearranging the code
-- so that an info table appears before its corresponding function.
--

module LlvmMangler ( llvmFixupAsm ) where

import DynFlags ( DynFlags )
import ErrUtils ( showPass )
import LlvmCodeGen.Ppr ( infoSection )

import Control.Exception
import Control.Monad ( when )
import qualified Data.ByteString.Char8 as B
import Data.Char
import System.IO

import Data.List ( sortBy )
import Data.Function ( on )

#if x86_64_TARGET_ARCH
#define REWRITE_AVX
#endif

-- Magic Strings
secStmt, infoSec, newLine, textStmt, dataStmt, syntaxUnified :: B.ByteString
secStmt       = B.pack "\t.section\t"
infoSec       = B.pack infoSection
newLine       = B.pack "\n"
textStmt      = B.pack "\t.text"
dataStmt      = B.pack "\t.data"
syntaxUnified = B.pack "\t.syntax unified"

infoLen :: Int
infoLen = B.length infoSec

-- Search Predicates
isType :: B.ByteString -> Bool
isType = B.isPrefixOf (B.pack "\t.type")

-- section of a file in the form of (header line, contents)
type Section = (B.ByteString, B.ByteString)

-- | Read in assembly file and process
llvmFixupAsm :: DynFlags -> FilePath -> FilePath -> IO ()
llvmFixupAsm dflags f1 f2 = {-# SCC "llvm_mangler" #-} do
    showPass dflags "LLVM Mangler"
    r <- openBinaryFile f1 ReadMode
    w <- openBinaryFile f2 WriteMode
    ss <- readSections r w
    hClose r
    let fixed = (map rewriteAVX . fixTables) ss
    mapM_ (writeSection w) fixed
    hClose w
    return ()

-- | This rewrites @.type@ annotations of function symbols to @%object@.
-- This is done as the linker can relocate @%functions@ through the
-- Procedure Linking Table (PLT). This is bad since we expect that the
-- info table will appear directly before the symbol's location. In the
-- case that the PLT is used, this will be not an info table but instead
-- some random PLT garbage.
rewriteSymType :: B.ByteString -> B.ByteString
rewriteSymType s =
    B.unlines $ map (rewrite '@' . rewrite '%') $ B.lines s
  where
    rewrite :: Char -> B.ByteString -> B.ByteString
    rewrite prefix x
        | isType x = replace funcType objType x
        | otherwise = x
      where
        funcType = prefix `B.cons` B.pack "function"
        objType  = prefix `B.cons` B.pack "object"

-- | Splits the file contents into its sections
readSections :: Handle -> Handle -> IO [Section]
readSections r w = go B.empty [] []
  where
    go hdr ss ls = do
      e_l <- (try (B.hGetLine r))::IO (Either IOError B.ByteString)

      -- Note that ".type" directives at the end of a section refer to
      -- the first directive of the *next* section, therefore we take
      -- it over to that section.
      let (tys, ls') = span isType ls
          cts = rewriteSymType $ B.intercalate newLine $ reverse ls'

      -- Decide whether to directly output the section or append it
      -- to the list for resorting.
      let finishSection
            | infoSec `B.isInfixOf` hdr =
                cts `seq` return $ (hdr, cts):ss
            | otherwise =
                writeSection w (hdr, cts) >> return ss

      case e_l of
        Right l | l == syntaxUnified
                  -> finishSection >>= \ss' -> writeSection w (l, B.empty)
                                   >> go B.empty ss' tys
                | any (`B.isPrefixOf` l) [secStmt, textStmt, dataStmt]
                  -> finishSection >>= \ss' -> go l ss' tys
                | otherwise
                  -> go hdr ss (l:ls)
        Left _    -> finishSection >>= \ss' -> return (reverse ss')

-- | Writes sections back
writeSection :: Handle -> Section -> IO ()
writeSection w (hdr, cts) = do
  when (not $ B.null hdr) $
    B.hPutStrLn w hdr
  B.hPutStrLn w cts

#if REWRITE_AVX
rewriteAVX :: Section -> Section
rewriteAVX = rewriteVmovaps . rewriteVmovdqa

rewriteVmovdqa :: Section -> Section
rewriteVmovdqa = rewriteInstructions vmovdqa vmovdqu
  where
    vmovdqa, vmovdqu :: B.ByteString
    vmovdqa = B.pack "vmovdqa"
    vmovdqu = B.pack "vmovdqu"

rewriteVmovap :: Section -> Section
rewriteVmovap = rewriteInstructions vmovap vmovup
  where
    vmovap, vmovup :: B.ByteString
    vmovap = B.pack "vmovap"
    vmovup = B.pack "vmovup"

rewriteInstructions :: B.ByteString -> B.ByteString -> Section -> Section
rewriteInstructions matchBS replaceBS (hdr, cts) =
    (hdr, replace matchBS replaceBS cts)
#else /* !REWRITE_AVX */
rewriteAVX :: Section -> Section
rewriteAVX = id
#endif /* !REWRITE_SSE */

replace :: B.ByteString -> B.ByteString -> B.ByteString -> B.ByteString
replace matchBS replaceBS = loop
  where
    loop :: B.ByteString -> B.ByteString
    loop cts =
        case B.breakSubstring matchBS cts of
          (hd,tl) | B.null tl -> hd
                  | otherwise -> hd `B.append` replaceBS `B.append`
                                 loop (B.drop (B.length matchBS) tl)

-- | Reorder and convert sections so info tables end up next to the
-- code. Also does stack fixups.
fixTables :: [Section] -> [Section]
fixTables ss = map strip sorted
  where
    -- Resort sections: We only assign a non-zero number to all
    -- sections having the "STRIP ME" marker. As sortBy is stable,
    -- this will cause all these sections to be appended to the end of
    -- the file in the order given by the indexes.
    extractIx hdr
      | B.null a  = 0
      | otherwise = 1 + readInt (B.takeWhile isDigit $ B.drop infoLen a)
      where (_,a) = B.breakSubstring infoSec hdr

    indexed = zip (map (extractIx . fst) ss) ss

    sorted = map snd $ sortBy (compare `on` fst) indexed

    -- Turn all the "STRIP ME" sections into normal text sections, as
    -- they are in the right place now.
    strip (hdr, cts)
      | infoSec `B.isInfixOf` hdr = (textStmt, cts)
      | otherwise                 = (hdr, cts)

-- | Read an int or error
readInt :: B.ByteString -> Int
readInt str | B.all isDigit str = (read . B.unpack) str
            | otherwise = error $ "LLvmMangler Cannot read " ++ show str
                                ++ " as it's not an Int"