Abstract the string table into a newtype

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,128 @@
+-- | Object file generation.
+module Language.Elna.Object.ElfCoder
+    ( ElfEnvironment(..)
+    , ElfWriter(..)
+    , ElfHeaderResult(..)
+    , elfHeaderSize
+    , addSectionHeader
+    , elfObject
+    , elfSectionsSize
+    , putSectionHeader
+    , module Language.Elna.Object.Elf
+    ) where
+
+import Control.Exception (throwIO)
+import Control.Monad.IO.Class (MonadIO(..))
+import Control.Monad.Trans.State (StateT, runStateT, modify', gets)
+import Data.ByteString (StrictByteString)
+import qualified Data.ByteString as ByteString
+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
+import GHC.Records (HasField(..))
+
+data ElfEnvironment = ElfEnvironment
+    { objectHeaders :: ElfHeaderResult Elf32_Shdr
+    , objectHandle :: Handle
+    }
+
+newtype ElfWriter a = ElfWriter
+    { runElfWriter :: StateT ElfEnvironment 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 elfEnvironment@ElfEnvironment{ objectHeaders } =
+        let ElfHeaderResult{..} = objectHeaders
+         in elfEnvironment
+            { objectHeaders = ElfHeaderResult
+                { sectionHeaders = Vector.snoc sectionHeaders newHeader
+                , sectionNames = StringTable.append name sectionNames
+                }
+            }
+
+putSectionHeader :: StrictByteString -> Elf32_Shdr -> StrictByteString -> ElfWriter ()
+putSectionHeader name newHeader encoded = do
+    objectHandle' <- ElfWriter $ gets $ getField @"objectHandle"
+    liftIO $ ByteString.hPut objectHandle' encoded
+    addSectionHeader name newHeader
+
+-- 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 -> ElfWriter Elf32_Ehdr -> IO ()
+elfObject outFile putContents = withFile outFile WriteMode withObjectFile
+  where
+    withObjectFile objectHandle
+        = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
+        >> putContents' objectHandle
+        >>= uncurry afterContents
+    putContents' objectHandle
+        = flip runStateT (initialState objectHandle)
+        $ runElfWriter putContents
+    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 objectHandle = ElfEnvironment
+        { objectHeaders = ElfHeaderResult
+            { sectionHeaders = Vector.singleton zeroHeader
+            , sectionNames = mempty
+            }
+        , objectHandle = objectHandle
+        }
+    afterContents header ElfEnvironment{ objectHeaders = 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