{-# 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"