Add a state monad transformer to the Elf generator

1 files changed, 30 insertions, 57 deletions

diff --git a/lib/Language/Elna/PrinterWriter.hs b/lib/Language/Elna/PrinterWriter.hs
index 95923cf..5575e16 100644
--- a/lib/Language/Elna/PrinterWriter.hs
+++ b/lib/Language/Elna/PrinterWriter.hs
@@ -35,51 +35,22 @@ import Language.Elna.Object.Elf
     , rInfo
     , elf32Rel
     , shfInfoLink
+    , ElfWriter(..)
+    , ElfHeaderResult(..)
+    , addSectionHeader
     )
 import System.IO (Handle)
 import qualified Language.Elna.Architecture.RiscV as RiscV
 import qualified Data.Text.Encoding as Text.Encoding
+import Control.Monad.IO.Class (MonadIO(..))
+import Control.Monad.Trans.State (get)
 
-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
+riscv32Elf :: Vector RiscV.Instruction -> Handle -> ElfWriter ()
+riscv32Elf code objectHandle = text
+    >>= symstrtab
   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
@@ -91,25 +62,26 @@ riscv32Elf code objectHandle =
                 , r_info = rInfo (fromIntegral entry) type'
                 }
         | otherwise = Left unresolvedRelocation
-    symstrtab (symbols@(ElfHeaderResult stringTable entries), ElfHeaderResult names headers, relocations) = do
+    symstrtab (symbols@(ElfHeaderResult stringTable entries), relocations) = do
+        ElfHeaderResult{..} <- ElfWriter get
         let encodedSymbols = LazyByteString.toStrict
                 $ ByteString.Builder.toLazyByteString
                 $ foldMap (elf32Sym LSB) entries 
-            namesLength = fromIntegral $ ByteString.length names
+            namesLength = fromIntegral $ ByteString.length sectionNames
             symHeader = Elf32_Shdr
                 { sh_type = SHT_SYMTAB
                 , sh_size = fromIntegral $ ByteString.length encodedSymbols
-                , sh_offset = elfSectionsSize headers
+                , sh_offset = elfSectionsSize sectionHeaders
                 , sh_name = namesLength
-                , sh_link = fromIntegral $ Vector.length headers + 2
+                , sh_link = fromIntegral $ Vector.length sectionHeaders + 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
+        liftIO $ ByteString.hPut objectHandle encodedSymbols
+        let headers1 = Vector.snoc sectionHeaders symHeader
         let y = resolveRelocation symbols <$> relocations
             encodedRelocations = LazyByteString.toStrict
                 $ ByteString.Builder.toLazyByteString
@@ -119,14 +91,14 @@ riscv32Elf code objectHandle =
                 , sh_size = fromIntegral $ ByteString.length encodedRelocations
                 , sh_offset = elfSectionsSize headers1
                 , sh_name = namesLength + 8
-                , sh_link = fromIntegral $ Vector.length headers
+                , sh_link = fromIntegral $ Vector.length sectionHeaders
                 , sh_info = 1
                 , sh_flags = shfInfoLink
                 , sh_entsize = 8
                 , sh_addralign = 0
                 , sh_addr = 0
                 }
-        ByteString.hPut objectHandle encodedRelocations
+        liftIO $ ByteString.hPut objectHandle encodedRelocations
         let headers2 = Vector.snoc headers1 relHeader
         let strHeader = Elf32_Shdr
                 { sh_type = SHT_STRTAB
@@ -140,11 +112,13 @@ riscv32Elf code objectHandle =
                 , 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
+        liftIO $ ByteString.hPut objectHandle stringTable
+        addSectionHeader ".symtab" symHeader
+        addSectionHeader ".rel.text" relHeader
+        addSectionHeader ".strtab" strHeader
+    text = do
+        ElfHeaderResult{..} <- ElfWriter get
+        let textTabIndex = fromIntegral $ Vector.length sectionHeaders
             initialHeaders = ElfHeaderResult "\0"
                 $ Vector.singleton
                 $ Elf32_Sym
@@ -160,8 +134,8 @@ riscv32Elf code objectHandle =
         let newHeader = Elf32_Shdr
                 { sh_type = SHT_PROGBITS
                 , sh_size = size
-                , sh_offset = elfSectionsSize headers
-                , sh_name = fromIntegral $ ByteString.length names
+                , sh_offset = elfSectionsSize sectionHeaders
+                , sh_name = fromIntegral $ ByteString.length sectionNames
                 , sh_link = 0
                 , sh_info = 0
                 , sh_flags = 0b110
@@ -169,14 +143,13 @@ riscv32Elf code objectHandle =
                 , sh_addralign = 0
                 , sh_addr = 0
                 }
-            newResult = ElfHeaderResult (names <> ".text\0")
-                $ Vector.snoc headers newHeader
-        pure (symbolResult, newResult, relocations)
+        addSectionHeader ".text" newHeader
+        pure (symbolResult, relocations)
     symbolEntry
         :: Elf32_Half
         -> Vector RiscV.Instruction
         -> (ElfHeaderResult Elf32_Sym, Elf32_Addr, Vector UnresolvedRelocation)
-        -> IO (ElfHeaderResult Elf32_Sym, Elf32_Addr, Vector UnresolvedRelocation)
+        -> ElfWriter (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
@@ -188,7 +161,7 @@ riscv32Elf code objectHandle =
                 , st_name = fromIntegral $ ByteString.length names
                 , st_info = stInfo STB_GLOBAL STT_FUNC
                 }
-        ByteString.hPut objectHandle $ LazyByteString.toStrict encoded
+        liftIO $ ByteString.hPut objectHandle $ LazyByteString.toStrict encoded
         let newResult = ElfHeaderResult (names <> "_start\0")
                 $ Vector.snoc entries newEntry
         pure (newResult, size, updatedRelocations)