Abstract the string table into a newtype

elna.cabal

@@ -4,12 +4,10 @@ version: 0.1.0.0
 synopsis:
     Elna programming language compiles simple mathematical operations to RISC-V code
 
--- description:
 license:            MPL-2.0
 license-file:       LICENSE
 author:             Eugen Wissner
 maintainer:         belka@caraus.de
--- copyright:
 category: Language
 build-type: Simple
 extra-doc-files: TODO README
@@ -48,6 +46,8 @@ library elna-internal
         Language.Elna.Glue
         Language.Elna.Location
         Language.Elna.Object.Elf
+        Language.Elna.Object.ElfCoder
+        Language.Elna.Object.StringTable
         Language.Elna.RiscV.CodeGenerator
         Language.Elna.RiscV.ElfWriter
     build-depends:

lib/Language/Elna/Object/Elf.hs

@@ -20,9 +20,7 @@ module Language.Elna.Object.Elf
     , ElfSectionType(..)
     , ElfSymbolBinding(..)
     , ElfSymbolType(..)
-    , ElfWriter(..)
-    , ElfHeaderResult(..)
-    , addSectionHeader
+    , byteOrder
     , elf32Addr
     , elf32Half
     , elf32Off
@@ -33,10 +31,7 @@ module Language.Elna.Object.Elf
     , elf32Rel
     , elf32Rela
     , elf32Sym
-    , elfHeaderSize
     , elfIdentification
-    , elfObject
-    , elfSectionsSize
     , rInfo
     , shfWrite
     , shfAlloc
@@ -46,19 +41,12 @@ module Language.Elna.Object.Elf
     , stInfo
     ) where
 
-import Control.Exception (Exception(..), throwIO)
+import Control.Exception (Exception(..))
 import Data.Bits (Bits(..))
 import qualified Data.ByteString.Builder as ByteString.Builder
 import Data.Int (Int32)
 import Data.Word (Word8, Word16, Word32)
 import qualified Data.ByteString as ByteString
-import Data.Vector (Vector)
-import qualified Data.Vector as Vector
-import System.IO (Handle, IOMode(..), SeekMode(..), hSeek, withFile)
-import Data.Foldable (traverse_)
-import Control.Monad.Trans.State (StateT, runStateT, modify')
-import Control.Monad.IO.Class (MonadIO(..))
-import Data.ByteString (ByteString)
 
 -- * Data types.
 
@@ -498,90 +486,3 @@ instance Exception ElfEncodingError
 
 fromIntegralEnum :: (Enum a, Num b) => a -> b
 fromIntegralEnum = fromIntegral . fromEnum
-
--- * Object file generation.
-
-newtype ElfWriter a = ElfWriter
-    { runElfWriter :: StateT (ElfHeaderResult Elf32_Shdr) IO a
-    }
-
-data ElfHeaderResult a = ElfHeaderResult
-    { sectionNames :: ByteString
-    , sectionHeaders :: Vector a
-    } deriving Eq
-
-instance Functor ElfWriter
-  where
-    fmap f (ElfWriter x) = ElfWriter $ f <$> x
-
-instance Applicative ElfWriter
-  where
-    pure = ElfWriter . pure
-    (ElfWriter f) <*> (ElfWriter x) = ElfWriter $ f <*> x
-
-instance Monad ElfWriter
-  where
-    (ElfWriter x) >>= f = ElfWriter $ x >>= (runElfWriter . f)
-
-instance MonadIO ElfWriter
-  where
-    liftIO = ElfWriter . liftIO
-
--- | ELF header size.
-elfHeaderSize :: Elf32_Off
-elfHeaderSize = 52
-
--- | Calculates the size of all sections based on the 'sh_size' in the given
--- headers and adds 'elfHeaderSize' to it.
-elfSectionsSize :: Vector Elf32_Shdr -> Elf32_Off
-elfSectionsSize = (elfHeaderSize +)
-    . Vector.foldr ((+) . sh_size) 0
-
-addSectionHeader :: ByteString -> Elf32_Shdr -> ElfWriter ()
-addSectionHeader name newHeader = ElfWriter $ modify' modifier
-  where
-    modifier ElfHeaderResult{..} =
-        ElfHeaderResult
-            { sectionHeaders = Vector.snoc sectionHeaders newHeader
-            , sectionNames = sectionNames <> name <> "\0"
-            }
-
--- Writes an ELF object to the provided file path. The callback writes the
--- sections, collects headers for those sections and returns the ELF header.
-elfObject :: FilePath -> (Handle -> ElfWriter Elf32_Ehdr) -> IO ()
-elfObject outFile putContents = withFile outFile WriteMode withObjectFile
-  where
-    withObjectFile objectHandle
-        = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
-        >> putContents' objectHandle
-        >>= uncurry (afterContents objectHandle)
-    putContents' objectHandle
-        = flip runStateT initialState
-        $ runElfWriter
-        $ putContents objectHandle
-    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
-        }
-    initialState = ElfHeaderResult
-        { sectionHeaders = Vector.singleton zeroHeader
-        , sectionNames = "\0"
-        }
-    afterContents objectHandle header ElfHeaderResult{..} =
-        let hPutBuilder = ByteString.Builder.hPutBuilder objectHandle
-            writeSectionHeaders byteOrder' =
-                traverse_ (hPutBuilder . elf32Shdr byteOrder') sectionHeaders
-         in either throwIO pure (byteOrder (e_ident header))
-            >>= writeSectionHeaders
-            >> either throwIO (putHeaders objectHandle) (elf32Ehdr header)
-    putHeaders objectHandle encodedHeader
-        = hSeek objectHandle AbsoluteSeek 0
-        >> ByteString.Builder.hPutBuilder objectHandle encodedHeader

lib/Language/Elna/Object/ElfCoder.hs

@@ -0,0 +1,108 @@
+-- | Object file generation.
+module Language.Elna.Object.ElfCoder
+    ( ElfWriter(..)
+    , ElfHeaderResult(..)
+    , elfHeaderSize
+    , addSectionHeader
+    , elfObject
+    , elfSectionsSize
+    , module Language.Elna.Object.Elf
+    ) where
+
+import Control.Exception (throwIO)
+import Control.Monad.IO.Class (MonadIO(..))
+import Control.Monad.Trans.State (StateT, runStateT, modify')
+import Data.ByteString (StrictByteString)
+import qualified Data.ByteString.Builder as ByteString.Builder
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import System.IO (Handle, IOMode(..), SeekMode(..), hSeek, withFile)
+import Data.Foldable (traverse_)
+import Language.Elna.Object.Elf
+import Language.Elna.Object.StringTable (StringTable)
+import qualified Language.Elna.Object.StringTable as StringTable
+
+newtype ElfWriter a = ElfWriter
+    { runElfWriter :: StateT (ElfHeaderResult Elf32_Shdr) IO a
+    }
+
+data ElfHeaderResult a = ElfHeaderResult
+    { sectionNames :: StringTable
+    , sectionHeaders :: Vector a
+    } deriving Eq
+
+instance Functor ElfWriter
+  where
+    fmap f (ElfWriter x) = ElfWriter $ f <$> x
+
+instance Applicative ElfWriter
+  where
+    pure = ElfWriter . pure
+    (ElfWriter f) <*> (ElfWriter x) = ElfWriter $ f <*> x
+
+instance Monad ElfWriter
+  where
+    (ElfWriter x) >>= f = ElfWriter $ x >>= (runElfWriter . f)
+
+instance MonadIO ElfWriter
+  where
+    liftIO = ElfWriter . liftIO
+
+-- | ELF header size.
+elfHeaderSize :: Elf32_Off
+elfHeaderSize = 52
+
+-- | Calculates the size of all sections based on the 'sh_size' in the given
+-- headers and adds 'elfHeaderSize' to it.
+elfSectionsSize :: Vector Elf32_Shdr -> Elf32_Off
+elfSectionsSize = (elfHeaderSize +)
+    . Vector.foldr ((+) . sh_size) 0
+
+addSectionHeader :: StrictByteString -> Elf32_Shdr -> ElfWriter ()
+addSectionHeader name newHeader = ElfWriter $ modify' modifier
+  where
+    modifier ElfHeaderResult{..} =
+        ElfHeaderResult
+            { sectionHeaders = Vector.snoc sectionHeaders newHeader
+            , sectionNames = StringTable.append name sectionNames
+            }
+
+-- Writes an ELF object to the provided file path. The callback writes the
+-- sections, collects headers for those sections and returns the ELF header.
+elfObject :: FilePath -> (Handle -> ElfWriter Elf32_Ehdr) -> IO ()
+elfObject outFile putContents = withFile outFile WriteMode withObjectFile
+  where
+    withObjectFile objectHandle
+        = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
+        >> putContents' objectHandle
+        >>= uncurry (afterContents objectHandle)
+    putContents' objectHandle
+        = flip runStateT initialState
+        $ runElfWriter
+        $ putContents objectHandle
+    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
+        }
+    initialState = ElfHeaderResult
+        { sectionHeaders = Vector.singleton zeroHeader
+        , sectionNames = mempty
+        }
+    afterContents objectHandle header ElfHeaderResult{..} =
+        let hPutBuilder = ByteString.Builder.hPutBuilder objectHandle
+            writeSectionHeaders byteOrder' =
+                traverse_ (hPutBuilder . elf32Shdr byteOrder') sectionHeaders
+         in either throwIO pure (byteOrder (e_ident header))
+            >>= writeSectionHeaders
+            >> either throwIO (putHeaders objectHandle) (elf32Ehdr header)
+    putHeaders objectHandle encodedHeader
+        = hSeek objectHandle AbsoluteSeek 0
+        >> ByteString.Builder.hPutBuilder objectHandle encodedHeader

lib/Language/Elna/Object/StringTable.hs

@@ -0,0 +1,44 @@
+module Language.Elna.Object.StringTable
+    ( StringTable
+    , append
+    , elem
+    , index
+    , encode
+    , size
+    ) where
+
+import Data.ByteString (StrictByteString)
+import qualified Data.ByteString as ByteString
+import Language.Elna.Object.Elf
+import Prelude hiding (elem)
+
+newtype StringTable = StringTable StrictByteString
+    deriving Eq
+
+instance Semigroup StringTable
+  where
+    (StringTable x) <> (StringTable y) = StringTable $ x <> ByteString.drop 1 y
+
+instance Monoid StringTable
+  where
+    mempty = StringTable "\0"
+
+size :: StringTable -> Elf32_Word
+size (StringTable container) =
+    fromIntegral $ ByteString.length container
+
+elem :: StrictByteString -> StringTable -> Bool
+elem needle (StringTable container) =
+    ("\0" <> needle <> "\0") `ByteString.isInfixOf` container
+
+append :: StrictByteString -> StringTable -> StringTable
+append element (StringTable container) =
+    StringTable $ container <> element <> "\0"
+
+index :: Elf32_Word -> StringTable -> StrictByteString
+index stringTableIndex (StringTable stringTable)
+    = ByteString.takeWhile (/= 0)
+    $ ByteString.drop (fromIntegral stringTableIndex) stringTable
+
+encode :: StringTable -> StrictByteString
+encode (StringTable container) = container

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(..)
@@ -43,12 +43,13 @@ 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
@@ -61,12 +62,14 @@ riscv32Elf code objectHandle = text
     shstrtab :: ElfWriter Elf32_Half
     shstrtab = do
         ElfHeaderResult{..} <- ElfWriter get
-        let stringTable = sectionNames <> ".shstrtab\0"
+        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,8 +77,10 @@ riscv32Elf code objectHandle = text
                 , sh_addralign = 1
                 , sh_addr = 0
                 }
-        liftIO $ ByteString.hPut objectHandle stringTable
-        addSectionHeader ".shstrtab" nextHeader
+        addSectionHeader stringTable nextHeader
+
+        updatedSectionNames <- ElfWriter $ gets (StringTable.encode . getField @"sectionNames")
+        liftIO $ ByteString.hPut objectHandle updatedSectionNames
         pure $ fromIntegral $ Vector.length sectionHeaders
     riscv32Header :: Elf32_Half -> ElfWriter Elf32_Ehdr
     riscv32Header shstrndx = do
@@ -97,8 +102,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 =
@@ -116,7 +120,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
@@ -142,7 +146,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
@@ -157,9 +161,9 @@ riscv32Elf code objectHandle = text
         ElfHeaderResult{..} <- ElfWriter get
         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 +171,12 @@ riscv32Elf code objectHandle = text
                 , sh_addralign = 1
                 , sh_addr = 0
                 }
-        liftIO $ ByteString.hPut objectHandle stringTable
+        liftIO $ ByteString.hPut objectHandle $ StringTable.encode stringTable
         addSectionHeader ".strtab" strHeader
     text = do
         ElfHeaderResult{..} <- ElfWriter get
         let textTabIndex = fromIntegral $ Vector.length sectionHeaders
-            initialHeaders = ElfHeaderResult "\0"
+            initialHeaders = ElfHeaderResult mempty
                 $ Vector.singleton
                 $ Elf32_Sym
                     { st_value = 0
@@ -185,9 +189,9 @@ riscv32Elf code objectHandle = text
             (encoded, updatedRelocations, symbols, definitions) = 
                 encodeAsm textTabIndex (mempty, Vector.empty, initialHeaders, HashSet.empty) code
 
+            filterPredicate :: StrictByteString -> Bool
             filterPredicate = not
-                . (`ByteString.isInfixOf` getField @"sectionNames" symbols)
-                . ("\0" <>) . (<> "\0")
+                . (`StringTable.elem` getField @"sectionNames" symbols)
             symbolResult = HashSet.foldl' encodeEmptyDefinitions symbols
                 $ HashSet.filter filterPredicate definitions
             size = fromIntegral $ LazyByteString.length encoded
@@ -195,7 +199,7 @@ riscv32Elf code objectHandle = text
                 { 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
@@ -212,10 +216,10 @@ 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
         | Just (instruction, rest) <- Vector.uncons instructions =
@@ -232,13 +236,14 @@ 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'

src/Main.hs

@@ -3,7 +3,7 @@ module Main
     ) where
 
 import Language.Elna.CommandLine (CommandLine(..), commandLine, execParser)
-import Language.Elna.Object.Elf (elfObject)
+import Language.Elna.Object.ElfCoder (elfObject)
 import Language.Elna.Backend.Allocator (allocate)
 import Language.Elna.Glue (glue)
 import Language.Elna.Frontend.NameAnalysis (nameAnalysis)