elna/lib/Language/Elna/PrinterWriter.hs

-- | Writer assembler to an object file.
module Language.Elna.PrinterWriter
    ( riscv32Elf
    ) where

import Data.Word (Word8)
import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Builder as ByteString.Builder
import qualified Data.ByteString.Lazy as LazyByteString
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Language.Elna.Object.Elf
    ( ByteOrder(..)
    , Elf32_Addr
    , Elf32_Ehdr(..)
    , Elf32_Half
    , Elf32_Word
    , Elf32_Sym(..)
    , ElfMachine(..)
    , ElfType(..)
    , ElfVersion(..)
    , ElfIdentification(..)
    , ElfClass(..)
    , ElfData(..)
    , Elf32_Shdr(..)
    , ElfSectionType(..)
    , ElfSymbolBinding(..)
    , ElfSymbolType(..)
    , Elf32_Rel(..)
    , ElfWriter(..)
    , ElfHeaderResult(..)
    , elf32Sym
    , elfHeaderSize
    , elfSectionsSize
    , stInfo
    , rInfo
    , elf32Rel
    , shfInfoLink
    , addSectionHeader
    )
import System.IO (Handle)
import qualified Language.Elna.Architecture.RiscV as RiscV
import qualified Data.Text.Encoding as Text.Encoding
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.State (get)

data UnresolvedRelocation = UnresolvedRelocation ByteString Elf32_Addr Word8
data UnresolvedRelocations = UnresolvedRelocations (Vector UnresolvedRelocation) Elf32_Word

riscv32Elf :: Vector RiscV.Instruction -> Handle -> ElfWriter Elf32_Ehdr
riscv32Elf code objectHandle = text
    >>= uncurry symrel
    >>= strtab
    >> shstrtab
    >>= riscv32Header
  where
    shstrtab :: ElfWriter Elf32_Half
    shstrtab = do
        ElfHeaderResult{..} <- ElfWriter get
        let stringTable = sectionNames <> ".shstrtab\0"
            nextHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = fromIntegral $ ByteString.length stringTable
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = fromIntegral $ ByteString.length sectionNames
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 1
                , sh_addr = 0
                }
        liftIO $ ByteString.hPut objectHandle stringTable
        addSectionHeader ".shstrtab" nextHeader
        pure $ fromIntegral $ Vector.length sectionHeaders
    riscv32Header :: Elf32_Half -> ElfWriter Elf32_Ehdr
    riscv32Header shstrndx = do
        ElfHeaderResult{..} <- ElfWriter get
        pure $ Elf32_Ehdr
            { e_version = EV_CURRENT
            , e_type = ET_REL
            , e_shstrndx = shstrndx
            , e_shoff = elfSectionsSize sectionHeaders
            , e_shnum = fromIntegral (Vector.length sectionHeaders)
            , e_shentsize = 40
            , e_phoff = 0
            , e_phnum = 0
            , e_phentsize = 32
            , e_machine = EM_RISCV
            , e_ident = ElfIdentification ELFCLASS32 ELFDATA2LSB
            , e_flags = 0x4 -- EF_RISCV_FLOAT_ABI_DOUBLE
            , e_entry = 0
            , e_ehsize = fromIntegral elfHeaderSize
            }
    takeStringZ stringTable Elf32_Sym{ st_name }
        = ByteString.takeWhile (/= 0)
        $ ByteString.drop (fromIntegral st_name) stringTable
    resolveRelocation (ElfHeaderResult stringTable entries) unresolvedRelocation
        | UnresolvedRelocation symbolName offset type' <- unresolvedRelocation
        , Just entry <- Vector.findIndex ((== symbolName) . takeStringZ stringTable) entries =
            Right $ Elf32_Rel
                { r_offset = offset
                , r_info = rInfo (fromIntegral entry) type'
                }
        | otherwise = Left unresolvedRelocation
    symtab entries = do 
        ElfHeaderResult{..} <- ElfWriter get
        let encodedSymbols = LazyByteString.toStrict
                $ ByteString.Builder.toLazyByteString
                $ foldMap (elf32Sym LSB) entries 
            symHeader = Elf32_Shdr
                { sh_type = SHT_SYMTAB
                , sh_size = fromIntegral $ ByteString.length encodedSymbols
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = fromIntegral $ ByteString.length sectionNames
                , sh_link = fromIntegral $ Vector.length sectionHeaders + 2
                , sh_info = 1
                , sh_flags = 0
                , sh_entsize = 16
                , sh_addralign = 4
                , sh_addr = 0
                }
        liftIO $ ByteString.hPut objectHandle encodedSymbols
        addSectionHeader ".symtab" symHeader
        pure $ fromIntegral $ Vector.length sectionHeaders
    symrel symbols relocations = do
        let UnresolvedRelocations relocationList index = relocations
            ElfHeaderResult stringTable entries = symbols

        sectionHeadersLength <- symtab entries
        ElfHeaderResult{..} <- ElfWriter get

        let encodedRelocations = LazyByteString.toStrict
                $ ByteString.Builder.toLazyByteString
                $ Vector.foldMap (either (const mempty) (elf32Rel LSB))
                $ resolveRelocation symbols <$> relocationList
            relHeader = Elf32_Shdr
                { sh_type = SHT_REL
                , sh_size = fromIntegral $ ByteString.length encodedRelocations
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = fromIntegral $ ByteString.length sectionNames
                , sh_link = sectionHeadersLength
                , sh_info = index
                , sh_flags = shfInfoLink
                , sh_entsize = 8
                , sh_addralign = 4
                , sh_addr = 0
                }
        liftIO $ ByteString.hPut objectHandle encodedRelocations
        addSectionHeader ".rel.text" relHeader
        pure stringTable
    strtab stringTable = do
        ElfHeaderResult{..} <- ElfWriter get
        let strHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = fromIntegral $ ByteString.length stringTable
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = fromIntegral $ ByteString.length sectionNames
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 1
                , sh_addr = 0
                }
        liftIO $ ByteString.hPut objectHandle stringTable
        addSectionHeader ".strtab" strHeader
    text = do
        ElfHeaderResult{..} <- ElfWriter get
        let textTabIndex = fromIntegral $ Vector.length sectionHeaders
            initialHeaders = ElfHeaderResult "\0"
                $ Vector.singleton
                $ Elf32_Sym
                    { st_value = 0
                    , st_size = 0
                    , st_shndx = 0
                    , st_other = 0
                    , st_name = 0
                    , st_info = 0
                    }
        (symbolResult, size, relocations) <- symbolEntry textTabIndex code
            (initialHeaders, 0, mempty)
        let newHeader = Elf32_Shdr
                { sh_type = SHT_PROGBITS
                , sh_size = size
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = fromIntegral $ ByteString.length sectionNames
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0b110
                , sh_entsize = 0
                , sh_addralign = 4
                , sh_addr = 0
                }
        addSectionHeader ".text" newHeader
        pure (symbolResult, UnresolvedRelocations relocations $ fromIntegral $ Vector.length sectionHeaders)
    symbolEntry
        :: Elf32_Half
        -> Vector RiscV.Instruction
        -> (ElfHeaderResult Elf32_Sym, Elf32_Addr, Vector UnresolvedRelocation)
        -> ElfWriter (ElfHeaderResult Elf32_Sym, Elf32_Addr, Vector UnresolvedRelocation)
    symbolEntry shndx instructions (ElfHeaderResult names entries, offset, relocations) = do
        let (encoded, size, updatedRelocations) = 
                Vector.foldl' encodeInstruction (mempty, offset, relocations) instructions
            newEntry = Elf32_Sym
                { st_value = offset
                , st_size = fromIntegral size
                , st_shndx = shndx
                , st_other = 0
                , st_name = fromIntegral $ ByteString.length names
                , st_info = stInfo STB_GLOBAL STT_FUNC
                }
            printEntry = Elf32_Sym
                { st_value = 0
                , st_size = 0
                , st_shndx = 0
                , st_other = 0
                , st_name = fromIntegral (ByteString.length names) + 7
                , st_info = stInfo STB_GLOBAL STT_FUNC
                }
        liftIO $ ByteString.hPut objectHandle $ LazyByteString.toStrict encoded
        let newResult = ElfHeaderResult (names <> "_start\0printi\0")
                $ Vector.snoc (Vector.snoc entries newEntry) printEntry
        pure (newResult, size, updatedRelocations)
    encodeInstruction (instructions, offset, relocations) instruction =
        let unresolvedRelocation = case instruction of
                RiscV.RelocatableInstruction _ instructionType
                    | RiscV.Higher20 _ symbolName <- instructionType
                        -> Just -- R_RISCV_HI20
                        $ UnresolvedRelocation (Text.Encoding.encodeUtf8 symbolName) offset 26
                    | RiscV.Lower12I _ _ _ symbolName <- instructionType
                        -> Just -- R_RISCV_LO12_I
                        $ UnresolvedRelocation (Text.Encoding.encodeUtf8 symbolName) offset 27
                    | RiscV.Lower12S symbolName _ _ _ <- instructionType
                        -> Just -- R_RISCV_LO12_S
                        $ UnresolvedRelocation (Text.Encoding.encodeUtf8 symbolName) offset 28
                RiscV.CallInstruction symbolName
                    -> Just -- R_RISCV_CALL_PLT
                    $ UnresolvedRelocation (Text.Encoding.encodeUtf8 symbolName) offset 19
                RiscV.Instruction _ _ -> Nothing
            encoded = ByteString.Builder.toLazyByteString
                $ RiscV.instruction instruction
         in
            ( instructions <> encoded
            , offset + fromIntegral (LazyByteString.length encoded)
            , maybe relocations (Vector.snoc relocations) unresolvedRelocation
            )