elna/lib/Language/Elna/PrinterWriter.hs

-- | Writer assembler to an object file.
module Language.Elna.PrinterWriter
    ( riscv32Elf
    , riscv32Header
    ) 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_Sym(..)
    , ElfMachine(..)
    , ElfType(..)
    , ElfVersion(..)
    , ElfIdentification(..)
    , ElfClass(..)
    , ElfData(..)
    , Elf32_Shdr(..)
    , ElfSectionType(..)
    , ElfSymbolBinding(..)
    , ElfSymbolType(..)
    , Elf32_Rel (..)
    , elf32Sym
    , elfHeaderSize
    , elfSectionsSize
    , stInfo
    , rInfo
    , elf32Rel
    )
import System.IO (Handle)
import qualified Language.Elna.Architecture.RiscV as RiscV
import qualified Data.Text.Encoding as Text.Encoding

data ElfHeaderResult a = ElfHeaderResult ByteString (Vector a)
data UnresolvedRelocation = UnresolvedRelocation ByteString Elf32_Addr Word8

riscv32Elf :: Vector RiscV.Instruction -> Handle -> IO (Vector Elf32_Shdr)
riscv32Elf code objectHandle =
    let 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
            }
     in text (ElfHeaderResult "\0" $ Vector.singleton zeroHeader)
        >>= symstrtab
        >>= shstrtab
        >>= finalize
  where
    finalize (ElfHeaderResult _ headers) = pure headers
    shstrtab (ElfHeaderResult names headers) = do
        let stringTable = names <> ".shstrtab\0"
            nextHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = fromIntegral $ ByteString.length stringTable
                , sh_offset = elfSectionsSize headers
                , sh_name = fromIntegral $ ByteString.length names
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 0
                , sh_addr = 0
                }
        ByteString.hPut objectHandle stringTable
        pure $ ElfHeaderResult stringTable
            $ Vector.snoc headers nextHeader
    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
    symstrtab (symbols@(ElfHeaderResult stringTable entries), ElfHeaderResult names headers, relocations) = do
        let encodedSymbols = LazyByteString.toStrict
                $ ByteString.Builder.toLazyByteString
                $ foldMap (elf32Sym LSB) entries 
            namesLength = fromIntegral $ ByteString.length names
            symHeader = Elf32_Shdr
                { sh_type = SHT_SYMTAB
                , sh_size = fromIntegral $ ByteString.length encodedSymbols
                , sh_offset = elfSectionsSize headers
                , sh_name = namesLength
                , sh_link = fromIntegral $ Vector.length headers + 2
                , sh_info = 1
                , sh_flags = 0
                , sh_entsize = 16
                , sh_addralign = 0
                , sh_addr = 0
                }
        ByteString.hPut objectHandle encodedSymbols
        let headers1 = Vector.snoc headers symHeader
        let y = resolveRelocation symbols <$> relocations
            encodedRelocations = LazyByteString.toStrict
                $ ByteString.Builder.toLazyByteString
                $ Vector.foldMap (either (const mempty) (elf32Rel LSB)) y
            relHeader = Elf32_Shdr
                { sh_type = SHT_REL
                , sh_size = fromIntegral $ ByteString.length encodedRelocations
                , sh_offset = elfSectionsSize headers1
                , sh_name = namesLength + 8
                , sh_link = fromIntegral $ Vector.length headers
                , sh_info = 1
                , sh_flags = 0x40 -- SHF_INFO_LINK
                , sh_entsize = 8
                , sh_addralign = 0
                , sh_addr = 0
                }
        ByteString.hPut objectHandle encodedRelocations
        let headers2 = Vector.snoc headers1 relHeader
        let strHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = fromIntegral $ ByteString.length stringTable
                , sh_offset = elfSectionsSize headers2
                , sh_name = namesLength + 18
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 0
                , sh_addr = 0
                }
        ByteString.hPut objectHandle stringTable
        pure $ ElfHeaderResult (names <> ".symtab\0.rel.text\0.strtab\0")
            $ Vector.snoc headers2 strHeader
    text (ElfHeaderResult names headers) = do
        let textTabIndex = fromIntegral $ Vector.length headers
            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 headers
                , sh_name = fromIntegral $ ByteString.length names
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0b110
                , sh_entsize = 0
                , sh_addralign = 0
                , sh_addr = 0
                }
            newResult = ElfHeaderResult (names <> ".text\0")
                $ Vector.snoc headers newHeader
        pure (symbolResult, newResult, relocations)
    symbolEntry
        :: Elf32_Half
        -> Vector RiscV.Instruction
        -> (ElfHeaderResult Elf32_Sym, Elf32_Addr, Vector UnresolvedRelocation)
        -> IO (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
                }
        ByteString.hPut objectHandle $ LazyByteString.toStrict encoded
        let newResult = ElfHeaderResult (names <> "_start\0")
                $ Vector.snoc entries newEntry
        pure (newResult, size, updatedRelocations)
    encodeInstruction (instructions, offset, relocations) instruction =
        let unresolvedRelocation = case instruction of
                RiscV.RelocatableInstruction _ instructionType
                    | RiscV.Higher20 _ symbolName <- instructionType
                        -> Just
                        $ UnresolvedRelocation (Text.Encoding.encodeUtf8 symbolName) offset 26
                    | RiscV.Lower12I _ _ _ symbolName <- instructionType
                        -> Just
                        $ UnresolvedRelocation (Text.Encoding.encodeUtf8 symbolName) offset 27
                    | RiscV.Lower12S symbolName _ _ _ <- instructionType
                        -> Just
                        $ UnresolvedRelocation (Text.Encoding.encodeUtf8 symbolName) offset 28
                RiscV.Instruction _ _ -> Nothing
            encoded = ByteString.Builder.toLazyByteString
                $ RiscV.instruction instruction
         in
            ( instructions <> encoded
            , offset + fromIntegral (LazyByteString.length encoded)
            , maybe relocations (Vector.snoc relocations) unresolvedRelocation
            )

riscv32Header :: Elf32_Ehdr
riscv32Header = Elf32_Ehdr
    { e_version = EV_CURRENT
    , e_type = ET_REL
    , e_shstrndx = 2 -- String table. SHN_UNDEF
    , e_shoff = 0
    , e_shnum = 0
    , 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
    }