Abstract the string table into a newtype

1 files changed, 51 insertions, 48 deletions

diff --git a/lib/Language/Elna/RiscV/ElfWriter.hs b/lib/Language/Elna/RiscV/ElfWriter.hs
index 52a92ea..d7723d3 100644
--- a/lib/Language/Elna/RiscV/ElfWriter.hs
+++ b/lib/Language/Elna/RiscV/ElfWriter.hs
@@ -4,13 +4,13 @@ module Language.Elna.RiscV.ElfWriter
     ) where
 
 import Data.Word (Word8)
-import Data.ByteString (ByteString)
+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.Elf
+import Language.Elna.Object.ElfCoder
     ( ByteOrder(..)
     , Elf32_Addr
     , Elf32_Ehdr(..)
@@ -30,6 +30,7 @@ import Language.Elna.Object.Elf
     , Elf32_Rel(..)
     , ElfWriter(..)
     , ElfHeaderResult(..)
+    , ElfEnvironment(..)
     , elf32Sym
     , elfHeaderSize
     , elfSectionsSize
@@ -38,21 +39,22 @@ import Language.Elna.Object.Elf
     , elf32Rel
     , shfInfoLink
     , addSectionHeader
+    , putSectionHeader
     )
-import System.IO (Handle)
 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)
+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 ByteString Elf32_Addr Word8
+data UnresolvedRelocation = UnresolvedRelocation StrictByteString Elf32_Addr Word8
 data UnresolvedRelocations = UnresolvedRelocations (Vector UnresolvedRelocation) Elf32_Word
 
-riscv32Elf :: Vector Statement -> Handle -> ElfWriter Elf32_Ehdr
-riscv32Elf code objectHandle = text
+riscv32Elf :: Vector Statement -> ElfWriter Elf32_Ehdr
+riscv32Elf code = text
     >>= uncurry symrel
     >>= strtab
     >> shstrtab
@@ -60,13 +62,15 @@ riscv32Elf code objectHandle = text
   where
     shstrtab :: ElfWriter Elf32_Half
     shstrtab = do
-        ElfHeaderResult{..} <- ElfWriter get
-        let stringTable = sectionNames <> ".shstrtab\0"
+        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
+        let stringTable = ".shstrtab"
+            currentNamesSize = StringTable.size sectionNames
             nextHeader = Elf32_Shdr
                 { sh_type = SHT_STRTAB
-                , sh_size = fromIntegral $ ByteString.length stringTable
+                , sh_size = currentNamesSize -- Adding trailing null character.
+                    + fromIntegral (succ $ ByteString.length stringTable)
                 , sh_offset = elfSectionsSize sectionHeaders
-                , sh_name = fromIntegral $ ByteString.length sectionNames
+                , sh_name = currentNamesSize
                 , sh_link = 0
                 , sh_info = 0
                 , sh_flags = 0
@@ -74,12 +78,16 @@ riscv32Elf code objectHandle = text
                 , sh_addralign = 1
                 , sh_addr = 0
                 }
-        liftIO $ ByteString.hPut objectHandle stringTable
-        addSectionHeader ".shstrtab" nextHeader
+        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 get
+        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
         pure $ Elf32_Ehdr
             { e_version = EV_CURRENT
             , e_type = ET_REL
@@ -97,8 +105,7 @@ riscv32Elf code objectHandle = text
             , e_ehsize = fromIntegral elfHeaderSize
             }
     takeStringZ stringTable Elf32_Sym{ st_name }
-        = ByteString.takeWhile (/= 0)
-        $ ByteString.drop (fromIntegral st_name) stringTable
+        = StringTable.index st_name stringTable
     resolveRelocation (ElfHeaderResult stringTable entries) unresolvedRelocation
         | UnresolvedRelocation symbolName offset type' <- unresolvedRelocation
         , Just entry <- Vector.findIndex ((== symbolName) . takeStringZ stringTable) entries =
@@ -108,7 +115,7 @@ riscv32Elf code objectHandle = text
                 }
         | otherwise = Left unresolvedRelocation
     symtab entries = do 
-        ElfHeaderResult{..} <- ElfWriter get
+        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
         let encodedSymbols = LazyByteString.toStrict
                 $ ByteString.Builder.toLazyByteString
                 $ foldMap (elf32Sym LSB) entries 
@@ -116,7 +123,7 @@ riscv32Elf code objectHandle = text
                 { sh_type = SHT_SYMTAB
                 , sh_size = fromIntegral $ ByteString.length encodedSymbols
                 , sh_offset = elfSectionsSize sectionHeaders
-                , sh_name = fromIntegral $ ByteString.length sectionNames
+                , sh_name = StringTable.size sectionNames
                 , sh_link = fromIntegral $ Vector.length sectionHeaders + 2
                 , sh_info = 1
                 , sh_flags = 0
@@ -124,15 +131,14 @@ riscv32Elf code objectHandle = text
                 , sh_addralign = 4
                 , sh_addr = 0
                 }
-        liftIO $ ByteString.hPut objectHandle encodedSymbols
-        addSectionHeader ".symtab" symHeader
+        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 get
+        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
 
         let encodedRelocations = LazyByteString.toStrict
                 $ ByteString.Builder.toLazyByteString
@@ -142,7 +148,7 @@ riscv32Elf code objectHandle = text
                 { sh_type = SHT_REL
                 , sh_size = fromIntegral $ ByteString.length encodedRelocations
                 , sh_offset = elfSectionsSize sectionHeaders
-                , sh_name = fromIntegral $ ByteString.length sectionNames
+                , sh_name = StringTable.size sectionNames
                 , sh_link = sectionHeadersLength
                 , sh_info = index
                 , sh_flags = shfInfoLink
@@ -150,16 +156,15 @@ riscv32Elf code objectHandle = text
                 , sh_addralign = 4
                 , sh_addr = 0
                 }
-        liftIO $ ByteString.hPut objectHandle encodedRelocations
-        addSectionHeader ".rel.text" relHeader
+        putSectionHeader ".rel.text" relHeader encodedRelocations
         pure stringTable
     strtab stringTable = do
-        ElfHeaderResult{..} <- ElfWriter get
+        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
         let strHeader = Elf32_Shdr
                 { sh_type = SHT_STRTAB
-                , sh_size = fromIntegral $ ByteString.length stringTable
+                , sh_size = StringTable.size stringTable
                 , sh_offset = elfSectionsSize sectionHeaders
-                , sh_name = fromIntegral $ ByteString.length sectionNames
+                , sh_name = StringTable.size sectionNames
                 , sh_link = 0
                 , sh_info = 0
                 , sh_flags = 0
@@ -167,12 +172,11 @@ riscv32Elf code objectHandle = text
                 , sh_addralign = 1
                 , sh_addr = 0
                 }
-        liftIO $ ByteString.hPut objectHandle stringTable
-        addSectionHeader ".strtab" strHeader
+        putSectionHeader ".strtab" strHeader $ StringTable.encode stringTable
     text = do
-        ElfHeaderResult{..} <- ElfWriter get
+        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
         let textTabIndex = fromIntegral $ Vector.length sectionHeaders
-            initialHeaders = ElfHeaderResult "\0"
+            initialHeaders = ElfHeaderResult mempty
                 $ Vector.singleton
                 $ Elf32_Sym
                     { st_value = 0
@@ -183,19 +187,13 @@ riscv32Elf code objectHandle = text
                     , st_info = 0
                     }
             (encoded, updatedRelocations, symbols, definitions) = 
-                encodeAsm textTabIndex (mempty, Vector.empty, initialHeaders, HashSet.empty) code
-
-            filterPredicate = not
-                . (`ByteString.isInfixOf` getField @"sectionNames" symbols)
-                . ("\0" <>) . (<> "\0")
-            symbolResult = HashSet.foldl' encodeEmptyDefinitions symbols
-                $ HashSet.filter filterPredicate 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 = fromIntegral $ ByteString.length sectionNames
+                , sh_name = StringTable.size sectionNames
                 , sh_link = 0
                 , sh_info = 0
                 , sh_flags = 0b110
@@ -203,8 +201,12 @@ riscv32Elf code objectHandle = text
                 , sh_addralign = 4
                 , sh_addr = 0
                 }
-        liftIO $ ByteString.hPut objectHandle $ LazyByteString.toStrict encoded
-        addSectionHeader ".text" newHeader
+        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
@@ -212,18 +214,18 @@ riscv32Elf code objectHandle = text
                 , st_size = 0
                 , st_shndx = 0
                 , st_other = 0
-                , st_name = fromIntegral (ByteString.length names)
+                , st_name = StringTable.size names
                 , st_info = stInfo STB_GLOBAL STT_FUNC
                 }
-         in ElfHeaderResult (names <> definition <> "\0")
+         in ElfHeaderResult (StringTable.append definition names)
             $ Vector.snoc entries nextEntry
-    encodeAsm shndx (encoded, relocations, ElfHeaderResult names symbols, definitions) instructions
+    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 encodeAsm shndx (encoded', relocations', ElfHeaderResult names symbols, definitions') rest'
+                     in encodeFunctions shndx (encoded', relocations', ElfHeaderResult names symbols, definitions') rest'
                 JumpLabel labelName _ ->
                     let (encoded', relocations', rest', definitions') =
                             encodeInstructions (encoded, relocations, rest, definitions)
@@ -232,16 +234,17 @@ riscv32Elf code objectHandle = text
                             , st_size = fromIntegral $ LazyByteString.length encoded'
                             , st_shndx = shndx
                             , st_other = 0
-                            , st_name = fromIntegral $ ByteString.length names
+                            , st_name = StringTable.size names
                             , st_info = stInfo STB_GLOBAL STT_FUNC
                             }
                         result =
                             ( encoded'
                             , relocations'
-                            , ElfHeaderResult (names <> Text.encodeUtf8 labelName <> "\0") (Vector.snoc symbols newEntry)
+                            , ElfHeaderResult (StringTable.append (Text.encodeUtf8 labelName) names)
+                                $ Vector.snoc symbols newEntry
                             , definitions'
                             )
-                     in encodeAsm shndx result rest'
+                     in encodeFunctions shndx result rest'
         | otherwise = (encoded, relocations, ElfHeaderResult names symbols, definitions)
     encodeInstructions (encoded, relocations, instructions, definitions)
         | Just (Instruction instruction, rest) <- Vector.uncons instructions =