Linker.hs

module Linker where

import Control.Monad.State

-- Relocation is either relative relocation (in current code block) or
-- symbol relocation -- a symbol from external library
type Relocation = Either Integer (String, String)
type Symbol = (String, Integer)

-- Whether a block of code is a library or not is stored in the memory
-- It's a hack to get mapping (LibraryName -> Address)
data CodeBlock = CodeBlock
  { cbSize :: Integer
  , cbLibraryName :: Maybe String
  , cbSymbols :: [Symbol]
  , cbRelocations :: [Relocation] }

data CodeBlockFinalized = CodeBlockFinalized
  { cbfSize :: Integer
  , cbfLibraryName :: Maybe String
  , cbfSymbols :: [Symbol]
  , cbfRelocations :: [Integer] }

data Elf = Elf
  { dependencies :: [String]
  , code :: CodeBlock }

data MemoryBlock = EmptyBlock Integer
                 | NonFinalizedBlock CodeBlock
                 | FinalizedBlock CodeBlockFinalized

type Memory = [MemoryBlock]

placeBlock :: Memory -> CodeBlock -> Memory
placeBlock [] _ = error "no space left! allocation aborted"
placeBlock (block@(EmptyBlock blockSize):rest) cb
  | cbSize cb == blockSize = NonFinalizedBlock cb : rest
  | cbSize cb < blockSize =
      NonFinalizedBlock cb : EmptyBlock (blockSize - cbSize cb) : rest
  | otherwise = block : placeBlock rest cb
placeBlock (block:rest) cb = block : placeBlock rest cb

placeBlockAt :: Memory -> CodeBlock -> Integer -> Memory
placeBlockAt [] _ _ = error "no space left! allocation aborted"
placeBlockAt (block@(EmptyBlock blockSize):rest) cb place
  | place < 0 = error "space is already occupied"
  | place < blockSize && place + cbSize cb == blockSize =
      EmptyBlock place : ncb : rest
  | place < blockSize && place + cbSize cb < blockSize =
      EmptyBlock place : ncb : EmptyBlock (blockSize - place - cbSize cb) : rest
  | place < blockSize && place + cbSize cb > blockSize =
      error "space is already occupied"
  | otherwise = placeBlockAt rest cb (place - blockSize)
  where ncb = NonFinalizedBlock cb
placeBlockAt (block@(NonFinalizedBlock code):rest) cb place =
  block : placeBlockAt rest cb (place - cbSize code)
placeBlockAt (block@(FinalizedBlock code):rest) cb place =
  block : placeBlockAt rest cb (place - cbfSize code)

type LinkerState = ([String], Memory, [(String, Elf)])

-- loadElf takes an executable in elf format and starting address
-- in memory
loadElf :: Elf -> Integer -> State LinkerState ()
loadElf elf offset = do
  (libsLoaded, memory, libs) <- get
  put (libsLoaded, placeBlockAt memory (code elf) offset, libs)
  loadDependencies $ dependencies elf
  assignRelocations

loadDependencies :: [String] -> State LinkerState ()
loadDependencies [] = return ()
loadDependencies (x:xs) = do
  (libsLoaded, memory, libs) <- get
  case x `elem` libsLoaded of
    True -> return ()
    False -> do
      let Just lib = lookup x libs
          memory = placeBlock memory (code lib)
          libsLoaded = x : libsLoaded
      put (libsLoaded, memory, libs)
      loadDependencies (dependencies lib)
      loadDependencies xs

librariesLocations :: Memory -> [(String, Integer, [Symbol])]
librariesLocations mem = helper 0 mem
  where helper start [] = []
        helper start ((EmptyBlock sz):rest) = helper (start + sz) rest
        helper start ((NonFinalizedBlock cb):rest) = case cbLibraryName cb of
          Nothing -> helper (start + cbSize cb) rest
          Just x -> (x, start, cbSymbols cb) : helper (start + cbSize cb) rest
        helper start ((FinalizedBlock c):rest) = helper (start + cbfSize c) rest

-- assignRelocations assumes that all dependencies were already loaded
assignRelocations :: State LinkerState ()
assignRelocations = do
  (libsLoaded, memory, libs) <- get
  let locations = librariesLocations memory
      memory = applyRelocations locations memory
  put (libsLoaded, memory, libs)

applyRelocations :: [(String, Integer, [(String, Integer)])] -> Memory -> Memory
applyRelocations locs mem = helper 0 mem
  where helper start [] = []
        helper start (b@(EmptyBlock sz):rest) = b : helper (start + sz) rest
        helper start (b@(FinalizedBlock c):rest) =
          b : helper (start + cbfSize c) rest
        helper start (b@(NonFinalizedBlock c):rest) =
          applyRelocation start c : helper (start + cbSize c) rest
        
        applyRelocation addr c = FinalizedBlock $ CodeBlockFinalized
                                 { cbfSize = cbSize c
                                 , cbfLibraryName = cbLibraryName c
                                 , cbfSymbols = cbSymbols c
                                 , cbfRelocations = map reloc $ cbRelocations c }
          where reloc :: Relocation -> Integer 
                reloc (Left offset) = addr + offset
                reloc (Right (lib, symbol)) = let
                  Just (libOffset, symbols) = lookup' locs lib
                  Just symbolOffset = lookup symbol symbols
                  in libOffset + symbolOffset

lookup' :: Eq a => [(a, b, c)] -> a -> Maybe (b, c)
lookup' [] _ = Nothing
lookup' ((a, b, c):rest) a'
  | a == a' = Just (b, c)
  | otherwise = lookup' rest a'