Add a state monad transformer to the Elf generator

lib/Language/Elna/Object/Elf.hs

@@ -20,6 +20,9 @@ module Language.Elna.Object.Elf
     , ElfSectionType(..)
     , ElfSymbolBinding(..)
     , ElfSymbolType(..)
+    , ElfWriter(..)
+    , ElfHeaderResult(..)
+    , addSectionHeader
     , elf32Addr
     , elf32Half
     , elf32Off
@@ -53,6 +56,9 @@ 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.
 
@@ -494,6 +500,32 @@ 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
@@ -504,26 +536,71 @@ 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 with the given header to the provided file path.
 -- The callback writes the sections and returns headers for those sections.
 --
 -- It updates some of the header header according to the given headers and
 -- expects .shstrtab be the last header in the list.
-elfObject :: FilePath -> Elf32_Ehdr -> (Handle -> IO (Vector Elf32_Shdr)) -> IO ()
+elfObject :: FilePath -> Elf32_Ehdr -> (Handle -> ElfWriter ()) -> IO ()
 elfObject outFile header putContents = withFile outFile WriteMode withObjectFile
   where
     withObjectFile objectHandle
         = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
-        >> putContents objectHandle
+        >> putContents' objectHandle
         >>= afterContents objectHandle
-    afterContents objectHandle headers =
-        let headerEncodingResult = elf32Ehdr
+    putContents' objectHandle
+        = fmap snd
+        $ 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 ElfHeaderResult{..} =
+        let stringTable = sectionNames <> ".shstrtab\0"
+            nextHeader = Elf32_Shdr
+                { sh_type = SHT_STRTAB
+                , sh_size = fromIntegral $ ByteString.length stringTable
+                , sh_offset = elfSectionsSize sectionHeaders
+                , sh_name = fromIntegral $ ByteString.length sectionNames
+                , sh_link = 0
+                , sh_info = 0
+                , sh_flags = 0
+                , sh_entsize = 0
+                , sh_addralign = 0
+                , sh_addr = 0
+                }
+            headers = Vector.snoc sectionHeaders nextHeader
+            headerEncodingResult = elf32Ehdr
                 $ header
                     { e_shoff = elfSectionsSize headers
                     , e_shnum = fromIntegral $ Vector.length headers
                     , e_shstrndx = fromIntegral (Vector.length headers) - 1
                     }
-         in traverse_ (ByteString.Builder.hPutBuilder objectHandle . elf32Shdr LSB) headers
+         in ByteString.hPut objectHandle stringTable
+            >> traverse_ (ByteString.Builder.hPutBuilder objectHandle . elf32Shdr LSB) headers
             >> either throwIO (putHeaders objectHandle) headerEncodingResult
     putHeaders objectHandle encodedHeader
         = hSeek objectHandle AbsoluteSeek 0