elna/lib/Language/Elna/RiscV/ElfWriter.hs

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

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_Ehdr(..)
    , Elf32_Half
    , Elf32_Word
    , Elf32_Sym(..)
    , ElfMachine(..)
    , ElfType(..)
    , ElfVersion(..)
    , ElfIdentification(..)
    , ElfClass(..)
    , ElfData(..)
    , Elf32_Shdr(..)
    , ElfSectionType(..)
    , ElfSymbolBinding(..)
    , ElfSymbolType(..)
    , Elf32_Rel(..)
    , ElfWriter(..)
    , ElfHeaderResult(..)
    , ElfEnvironment(..)
    , UnresolvedRelocation(..)
    , UnresolvedRelocations(..)
    , addHeaderToResult
    , addSectionHeader
    , elf32Sym
    , elfHeaderSize
    , elfSectionsSize
    , stInfo
    , rInfo
    , elf32Rel
    , shfInfoLink
    , partitionSymbols
    , 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 (Directive(..), Statement(..))
import Language.Elna.Object.StringTable (StringTable)
import qualified Language.Elna.Object.StringTable as StringTable
import qualified Data.HashSet as HashSet
import GHC.Records (HasField(..))

riscv32Elf :: Vector Statement -> ElfWriter Elf32_Ehdr
riscv32Elf code = text code
    >>= symtab
    >>= uncurry symrel
    >>= strtab
    >> shstrtab
    >>= riscv32Header
  where
    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
            }

text :: Vector Statement -> ElfWriter UnresolvedRelocations
text code = 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 $ UnresolvedRelocations updatedRelocations symbolResult
        $ fromIntegral $ Vector.length sectionHeaders
  where
    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, symbolResult, definitions) instructions
        | Just (instruction, rest) <- Vector.uncons instructions =
            case instruction of
                Instruction _ ->
                    let (encoded', relocations', symbolResult', definitions', rest') =
                            encodeInstructions shndx (encoded, relocations, symbolResult, definitions, instructions)
                     in encodeFunctions shndx (encoded', relocations', symbolResult', definitions') rest'
                JumpLabel labelName directives ->
                    let (encoded', relocations', ElfHeaderResult _names _symbols, definitions', rest') =
                            encodeInstructions shndx (encoded, relocations, symbolResult, definitions, rest)
                        isGlobal = GlobalDirective `elem` directives
                        newEntry = Elf32_Sym
                            { st_value = fromIntegral $ LazyByteString.length encoded
                            , st_size = if isGlobal then fromIntegral $ LazyByteString.length encoded' else 0
                            , st_shndx = shndx
                            , st_other = 0
                            , st_name = StringTable.size _names
                            , st_info = stInfo (if isGlobal then STB_GLOBAL else STB_LOCAL)
                                $ if FunctionDirective `elem` directives then STT_FUNC else STT_NOTYPE
                            }
                        result =
                            ( encoded'
                            , relocations'
                            , ElfHeaderResult (StringTable.append (Text.encodeUtf8 labelName) _names)
                                $ Vector.snoc _symbols newEntry
                            , definitions'
                            )
                     in encodeFunctions shndx result rest'
        | otherwise = (encoded, relocations, symbolResult, definitions)
    encodeInstructions shndx (encoded, relocations, ElfHeaderResult names symbols, definitions, instructions)
        | 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
                    , ElfHeaderResult names symbols
                    , addDefinition unresolvedRelocation definitions
                    , rest
                    )
             in encodeInstructions shndx result
        | Just (JumpLabel labelName directives , rest) <- Vector.uncons instructions
        , FunctionDirective `notElem` directives =
            let newEntry = Elf32_Sym
                    { st_value = fromIntegral $ LazyByteString.length encoded
                    , st_size = 0
                    , st_shndx = shndx
                    , st_other = 0
                    , st_name = StringTable.size names
                    , st_info = stInfo (if GlobalDirective `elem` directives then STB_GLOBAL else STB_LOCAL) STT_NOTYPE
                    }
                result =
                    ( encoded
                    , relocations
                    , ElfHeaderResult (StringTable.append (Text.encodeUtf8 labelName) names)
                        $ Vector.snoc symbols newEntry
                    , definitions
                    , rest
                    )
                in encodeInstructions shndx result
        | otherwise = (encoded, relocations, ElfHeaderResult names symbols, definitions, instructions)
    addDefinition (Just (UnresolvedRelocation symbolName _ _)) =
        HashSet.insert symbolName
    addDefinition Nothing = id

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

symtab :: UnresolvedRelocations -> ElfWriter (Elf32_Word, UnresolvedRelocations)
symtab (UnresolvedRelocations relocationList symbolResult index) = do 
    ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
    let (localSymbols, globalSymbols) = partitionSymbols symbolResult
        sortedSymbols = localSymbols <> globalSymbols
        sortedResult = symbolResult{ sectionHeaders = sortedSymbols }
        encodedSymbols = LazyByteString.toStrict
            $ ByteString.Builder.toLazyByteString
            $ foldMap (elf32Sym LSB) sortedSymbols
        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 = fromIntegral $ Vector.length localSymbols
            , sh_flags = 0
            , sh_entsize = 16
            , sh_addralign = 4
            , sh_addr = 0
            }
    putSectionHeader ".symtab" symHeader encodedSymbols
    pure
        ( fromIntegral $ Vector.length sectionHeaders
        , UnresolvedRelocations relocationList sortedResult index
        )

symrel :: Elf32_Word -> UnresolvedRelocations -> ElfWriter StringTable
symrel sectionHeadersLength relocations = do
    ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"

    let UnresolvedRelocations relocationList symbols index = relocations
        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 $ getField @"sectionNames" symbols
  where
    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

strtab :: StringTable -> ElfWriter ()
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