-- | Object file generation.
module Language.Elna.Object.ElfCoder
    ( ElfWriter(..)
    , ElfHeaderResult(..)
    , elfHeaderSize
    , addSectionHeader
    , elfObject
    , elfSectionsSize
    , module Language.Elna.Object.Elf
    ) where

import Control.Exception (throwIO)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.State (StateT, runStateT, modify')
import Data.ByteString (StrictByteString)
import qualified Data.ByteString.Builder as ByteString.Builder
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import System.IO (Handle, IOMode(..), SeekMode(..), hSeek, withFile)
import Data.Foldable (traverse_)
import Language.Elna.Object.Elf
import Language.Elna.Object.StringTable (StringTable)
import qualified Language.Elna.Object.StringTable as StringTable

newtype ElfWriter a = ElfWriter
    { runElfWriter :: StateT (ElfHeaderResult Elf32_Shdr) IO a
    }

data ElfHeaderResult a = ElfHeaderResult
    { sectionNames :: StringTable
    , sectionHeaders :: Vector a
    } deriving Eq

instance Functor ElfWriter
  where
    fmap f (ElfWriter x) = ElfWriter $ f <$> x

instance Applicative ElfWriter
  where
    pure = ElfWriter . pure
    (ElfWriter f) <*> (ElfWriter x) = ElfWriter $ f <*> x

instance Monad ElfWriter
  where
    (ElfWriter x) >>= f = ElfWriter $ x >>= (runElfWriter . f)

instance MonadIO ElfWriter
  where
    liftIO = ElfWriter . liftIO

-- | ELF header size.
elfHeaderSize :: Elf32_Off
elfHeaderSize = 52

-- | Calculates the size of all sections based on the 'sh_size' in the given
-- headers and adds 'elfHeaderSize' to it.
elfSectionsSize :: Vector Elf32_Shdr -> Elf32_Off
elfSectionsSize = (elfHeaderSize +)
    . Vector.foldr ((+) . sh_size) 0

addSectionHeader :: StrictByteString -> Elf32_Shdr -> ElfWriter ()
addSectionHeader name newHeader = ElfWriter $ modify' modifier
  where
    modifier ElfHeaderResult{..} =
        ElfHeaderResult
            { sectionHeaders = Vector.snoc sectionHeaders newHeader
            , sectionNames = StringTable.append name sectionNames
            }

-- Writes an ELF object to the provided file path. The callback writes the
-- sections, collects headers for those sections and returns the ELF header.
elfObject :: FilePath -> (Handle -> ElfWriter Elf32_Ehdr) -> IO ()
elfObject outFile putContents = withFile outFile WriteMode withObjectFile
  where
    withObjectFile objectHandle
        = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
        >> putContents' objectHandle
        >>= uncurry (afterContents objectHandle)
    putContents' objectHandle
        = flip runStateT initialState
        $ runElfWriter
        $ putContents objectHandle
    zeroHeader = Elf32_Shdr
        { sh_type = SHT_NULL
        , sh_size = 0
        , sh_offset = 0
        , sh_name = 0
        , sh_link = 0
        , sh_info = 0
        , sh_flags = 0
        , sh_entsize = 0
        , sh_addralign = 0
        , sh_addr = 0
        }
    initialState = ElfHeaderResult
        { sectionHeaders = Vector.singleton zeroHeader
        , sectionNames = mempty
        }
    afterContents objectHandle header ElfHeaderResult{..} =
        let hPutBuilder = ByteString.Builder.hPutBuilder objectHandle
            writeSectionHeaders byteOrder' =
                traverse_ (hPutBuilder . elf32Shdr byteOrder') sectionHeaders
         in either throwIO pure (byteOrder (e_ident header))
            >>= writeSectionHeaders
            >> either throwIO (putHeaders objectHandle) (elf32Ehdr header)
    putHeaders objectHandle encodedHeader
        = hSeek objectHandle AbsoluteSeek 0
        >> ByteString.Builder.hPutBuilder objectHandle encodedHeader