path: root/lib/Language/Elna/RiscV/ElfWriter.hs

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

import Data.Word (Word8)
import Data.ByteString (StrictByteString)
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.ElfCoder
    ( 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(..)
    , ElfEnvironment(..)
    , elf32Sym
    , elfHeaderSize
    , elfSectionsSize
    , stInfo
    , rInfo
    , elf32Rel
    , shfInfoLink
    , addSectionHeader
    , putSectionHeader
    )
import qualified Language.Elna.Architecture.RiscV as RiscV
import qualified Data.Text.Encoding as Text
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.State (get, gets)
import Language.Elna.RiscV.CodeGenerator (Statement(..))
import qualified Language.Elna.Object.StringTable as StringTable
import qualified Data.HashSet as HashSet
import GHC.Records (HasField(..))

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

riscv32Elf :: Vector Statement -> ElfWriter Elf32_Ehdr
riscv32Elf code = text
    >>= uncurry symrel
    >>= strtab
    >> shstrtab
    >>= riscv32Header
  where
    shstrtab :: ElfWriter Elf32_Half
    shstrtab = do
        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
        let stringTable = ".shstrtab"
            currentNamesSize = StringTable.size sectionNames
            nextHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = currentNamesSize -- Adding trailing null character.
                    + fromIntegral (succ $ ByteString.length stringTable)
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = currentNamesSize
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 1
                , sh_addr = 0
                }
        addSectionHeader stringTable nextHeader

        ElfEnvironment{..} <- ElfWriter get
        liftIO $ ByteString.hPut objectHandle
            $ StringTable.encode
            $ getField @"sectionNames" objectHeaders
        pure $ fromIntegral $ Vector.length sectionHeaders
    riscv32Header :: Elf32_Half -> ElfWriter Elf32_Ehdr
    riscv32Header shstrndx = do
        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
        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 }
        = StringTable.index 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 $ gets $ getField @"objectHeaders"
        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 = StringTable.size sectionNames
                , sh_link = fromIntegral $ Vector.length sectionHeaders + 2
                , sh_info = 1
                , sh_flags = 0
                , sh_entsize = 16
                , sh_addralign = 4
                , sh_addr = 0
                }
        putSectionHeader ".symtab" symHeader encodedSymbols
        pure $ fromIntegral $ Vector.length sectionHeaders
    symrel symbols relocations = do
        let UnresolvedRelocations relocationList index = relocations
            ElfHeaderResult stringTable entries = symbols

        sectionHeadersLength <- symtab entries
        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"

        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 = StringTable.size sectionNames
                , sh_link = sectionHeadersLength
                , sh_info = index
                , sh_flags = shfInfoLink
                , sh_entsize = 8
                , sh_addralign = 4
                , sh_addr = 0
                }
        putSectionHeader ".rel.text" relHeader encodedRelocations
        pure stringTable
    strtab stringTable = do
        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
        let strHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = StringTable.size stringTable
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = StringTable.size sectionNames
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 1
                , sh_addr = 0
                }
        putSectionHeader ".strtab" strHeader $ StringTable.encode stringTable
    text = do
        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
        let textTabIndex = fromIntegral $ Vector.length sectionHeaders
            initialHeaders = ElfHeaderResult mempty
                $ Vector.singleton
                $ Elf32_Sym
                    { st_value = 0
                    , st_size = 0
                    , st_shndx = 0
                    , st_other = 0
                    , st_name = 0
                    , st_info = 0
                    }
            (encoded, updatedRelocations, symbols, definitions) = 
                encodeFunctions textTabIndex (mempty, Vector.empty, initialHeaders, HashSet.empty) code
            size = fromIntegral $ LazyByteString.length encoded
            newHeader = Elf32_Shdr
                { sh_type = SHT_PROGBITS
                , sh_size = size
                , sh_offset = elfSectionsSize sectionHeaders
                , sh_name = StringTable.size sectionNames
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0b110
                , sh_entsize = 0
                , sh_addralign = 4
                , sh_addr = 0
                }
        putSectionHeader ".text" newHeader $ LazyByteString.toStrict encoded
        let filterPredicate :: StrictByteString -> Bool
            filterPredicate = not
                . (`StringTable.elem` getField @"sectionNames" symbols)
            symbolResult = HashSet.foldl' encodeEmptyDefinitions symbols
                $ HashSet.filter filterPredicate definitions
        pure (symbolResult, UnresolvedRelocations updatedRelocations $ fromIntegral $ Vector.length sectionHeaders)
    encodeEmptyDefinitions (ElfHeaderResult names entries) definition =
        let nextEntry = Elf32_Sym
                { st_value = 0
                , st_size = 0
                , st_shndx = 0
                , st_other = 0
                , st_name = StringTable.size names
                , st_info = stInfo STB_GLOBAL STT_FUNC
                }
         in ElfHeaderResult (StringTable.append definition names)
            $ Vector.snoc entries nextEntry
    encodeFunctions shndx (encoded, relocations, ElfHeaderResult names symbols, definitions) instructions
        | Just (instruction, rest) <- Vector.uncons instructions =
            case instruction of
                Instruction _ ->
                    let (encoded', relocations', rest', definitions') =
                            encodeInstructions (encoded, relocations, instructions, definitions)
                     in encodeFunctions shndx (encoded', relocations', ElfHeaderResult names symbols, definitions') rest'
                JumpLabel labelName _ ->
                    let (encoded', relocations', rest', definitions') =
                            encodeInstructions (encoded, relocations, rest, definitions)
                        newEntry = Elf32_Sym
                            { st_value = fromIntegral $ LazyByteString.length encoded
                            , st_size = fromIntegral $ LazyByteString.length encoded'
                            , st_shndx = shndx
                            , st_other = 0
                            , st_name = StringTable.size names
                            , st_info = stInfo STB_GLOBAL STT_FUNC
                            }
                        result =
                            ( encoded'
                            , relocations'
                            , ElfHeaderResult (StringTable.append (Text.encodeUtf8 labelName) names)
                                $ Vector.snoc symbols newEntry
                            , definitions'
                            )
                     in encodeFunctions shndx result rest'
        | otherwise = (encoded, relocations, ElfHeaderResult names symbols, definitions)
    encodeInstructions (encoded, relocations, instructions, definitions)
        | Just (Instruction instruction, rest) <- Vector.uncons instructions =
            let offset = fromIntegral $ LazyByteString.length encoded
                unresolvedRelocation = case instruction of
                    RiscV.RelocatableInstruction _ instructionType
                        | RiscV.RHigher20 _ symbolName <- instructionType
                            -> Just -- R_RISCV_HI20
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 26
                        | RiscV.RLower12I _ _ _ symbolName <- instructionType
                            -> Just -- R_RISCV_LO12_I
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 27
                        | RiscV.RLower12S symbolName _ _ _ <- instructionType
                            -> Just -- R_RISCV_LO12_S
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 28
                        | RiscV.RBranch symbolName _ _ _ <- instructionType
                            -> Just -- R_RISCV_BRANCH
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 16
                        | RiscV.RJal _ symbolName <- instructionType
                            -> Just -- R_RISCV_JAL
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 17
                    RiscV.CallInstruction symbolName
                        -> Just -- R_RISCV_CALL_PLT
                        $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 19
                    RiscV.BaseInstruction _ _ -> Nothing
                chunk = ByteString.Builder.toLazyByteString
                    $ RiscV.instruction instruction
                result =
                    ( encoded <> chunk
                    , maybe relocations (Vector.snoc relocations) unresolvedRelocation
                    , rest
                    , addDefinition unresolvedRelocation definitions
                    )
             in encodeInstructions result
        | otherwise = (encoded, relocations, instructions, definitions)
    addDefinition (Just (UnresolvedRelocation symbolName _ _)) =
        HashSet.insert symbolName
    addDefinition Nothing = id