melf-1.0.2: examples/AsmAArch64.hs
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}
module AsmAArch64
( CodeState
, Register
, Label
, adr
, b
, mov
, ldr
, svc
, ascii
, label
, exportSymbol
, assemble
, x0, x1, x2, x8
, w0, w1
) where
import Prelude as P
import Control.Exception.ChainedException
import Control.Monad.Catch
import Control.Monad.State as MS
import Data.Bits
import Data.ByteString.Builder
import Data.ByteString.Lazy as BSL
import Data.ByteString.Lazy.Char8 as BSLC
import Data.Int
import Data.Kind
import Data.Singletons.Sigma
import Data.Singletons.TH
import Data.Word
import Data.Elf
import Data.Elf.Constants
import Data.Elf.Headers
$(singletons [d| data RegisterWidth = X | W |])
type Register :: RegisterWidth -> Type
newtype Register c = R Word32
newtype CodeOffset = CodeOffset { getCodeOffset :: Int64 } deriving (Eq, Show, Ord, Num, Enum, Real, Integral, Bits, FiniteBits)
newtype Instruction = Instruction { getInstruction :: Word32 } deriving (Eq, Show, Ord, Num, Enum, Real, Integral, Bits, FiniteBits)
data Label = CodeRef !CodeOffset
| PoolRef !CodeOffset
-- Args:
-- Offset of the instruction
-- Offset of the pool
type InstructionGen = CodeOffset -> CodeOffset -> Either String Instruction
data CodeState = CodeState
{ offsetInPool :: CodeOffset
, poolReversed :: [Builder]
, codeReversed :: [InstructionGen]
, symbolsRefersed :: [(String, Label)]
}
emit' :: MonadState CodeState m => InstructionGen -> m ()
emit' g = modify f where
f CodeState {..} = CodeState { codeReversed = g : codeReversed
, ..
}
emit :: MonadState CodeState m => Instruction -> m ()
emit i = emit' $ \ _ _ -> Right i
emitPool :: MonadState CodeState m => Word -> ByteString -> m Label
emitPool a bs = state f where
f CodeState {..} =
let
offsetInPool' = align a offsetInPool
o = builderRepeatZero $ fromIntegral $ offsetInPool' - offsetInPool
in
( PoolRef offsetInPool'
, CodeState { offsetInPool = fromIntegral (BSL.length bs) + offsetInPool'
, poolReversed = lazyByteString bs : o : poolReversed
, ..
}
)
label :: MonadState CodeState m => m Label
label = gets (CodeRef . (* instructionSize) . fromIntegral . P.length . codeReversed)
x0, x1, x2, x8 :: Register 'X
x0 = R 0
x1 = R 1
x2 = R 2
x8 = R 8
w0, w1 :: Register 'W
w0 = R 0
w1 = R 1
isPower2 :: (Bits i, Num i) => i -> Bool
isPower2 n = n .&. (n - 1) == 0
align :: Word -> CodeOffset -> CodeOffset
align a _ | not (isPower2 a) = error "align is not a power of 2"
align 0 n = n
align a n = (n + a' - 1) .&. complement (a' - 1)
where a' = fromIntegral a
builderRepeatZero :: Int -> Builder
builderRepeatZero n = mconcat $ P.replicate n (word8 0)
b64 :: forall w . SingI w => Register w -> Word32
b64 _ = case sing @ w of
SX -> 1
SW -> 0
-- | C6.2.10 ADR
adr :: MonadState CodeState m => Register 'X -> Label -> m ()
adr (R n) rr = emit' f where
offsetToImm :: CodeOffset -> Either String Word32
offsetToImm (CodeOffset o) =
if not $ isBitN 19 o
then Left "offset is too big"
else
let
immlo = o .&. 3
immhi = (o `shiftR` 2) .&. 0x7ffff
in
Right $ fromIntegral $ (immhi `shift` 5) .|. (immlo `shift` 29)
f :: InstructionGen
f instrAddr poolOffset = do
imm <- offsetToImm $ findOffset poolOffset rr - instrAddr
return $ Instruction $ 0x10000000
.|. imm
.|. n
-- | C6.2.26 B
b :: MonadState CodeState m => Label -> m ()
b rr = emit' f where
offsetToImm26 :: CodeOffset -> Either String Word32
offsetToImm26 (CodeOffset o)
| o .&. 0x3 /= 0 = Left $ "offset is not aligned: " ++ show o
| not $ isBitN 28 o = Left "offset is too big"
| otherwise = Right $ fromIntegral $ o `shiftR` 2
f :: InstructionGen
f instrAddr poolOffset = do
imm26 <- offsetToImm26 $ findOffset poolOffset rr - instrAddr
return $ Instruction $ 0x14000000 .|. imm26
-- | C6.2.187 MOV (wide immediate)
mov :: (MonadState CodeState m, SingI w) => Register w -> Word16 -> m ()
mov r@(R n) imm = emit $ Instruction $ (b64 r `shift` 31)
.|. 0x52800000
.|. (fromIntegral imm `shift` 5)
.|. n
-- | The number can be represented with bitN bits
isBitN ::(Num b, Bits b, Ord b) => Int -> b -> Bool
isBitN bitN w =
let
m = complement $ (1 `shift` bitN) - 1
h = w .&. m
in if w >= 0 then h == 0 else h == m
findOffset :: CodeOffset -> Label -> CodeOffset
findOffset _poolOffset (CodeRef codeOffset) = codeOffset
findOffset poolOffset (PoolRef offsetInPool) = poolOffset + offsetInPool
-- | C6.2.132 LDR (literal)
ldr :: (MonadState CodeState m, SingI w) => Register w -> Label -> m ()
ldr r@(R n) rr = emit' f where
offsetToImm19 :: CodeOffset -> Either String Word32
offsetToImm19 (CodeOffset o)
| o .&. 0x3 /= 0 = Left "offset is not aligned"
| not $ isBitN 21 o = Left "offset is too big"
| otherwise = Right $ fromIntegral $ o `shiftR` 2
f :: InstructionGen
f instrAddr poolOffset = do
imm19 <- offsetToImm19 $ findOffset poolOffset rr - instrAddr
return $ Instruction $ (b64 r `shift` 30)
.|. 0x18000000
.|. (imm19 `shift` 5)
.|. n
-- | C6.2.317 SVC
svc :: MonadState CodeState m => Word16 -> m ()
svc imm = emit $ 0xd4000001 .|. (fromIntegral imm `shift` 5)
ascii :: MonadState CodeState m => String -> m Label
ascii s = emitPool 1 $ BSLC.pack s
exportSymbol :: MonadState CodeState m => String -> Label -> m ()
exportSymbol s r = modify f where
f (CodeState {..}) = CodeState { symbolsRefersed = (s, r) : symbolsRefersed
, ..
}
instructionSize :: Num b => b
instructionSize = 4
zeroIndexStringItem :: ElfSymbolXX 'ELFCLASS64
zeroIndexStringItem = ElfSymbolXX "" 0 0 0 0 0
textSecN, shstrtabSecN, strtabSecN, symtabSecN :: ElfSectionIndex
textSecN = 1
shstrtabSecN = 2
strtabSecN = 3
symtabSecN = 4
assemble :: MonadCatch m => StateT CodeState m () -> m Elf
assemble m = do
CodeState {..} <- execStateT m (CodeState 0 [] [] [])
-- resolve txt
let
poolOffset = instructionSize * fromIntegral (P.length codeReversed)
poolOffsetAligned = align 8 poolOffset
f :: (InstructionGen, CodeOffset) -> Either String Instruction
f (ff, n) = ff n poolOffsetAligned
code <- $eitherAddContext' $ mapM f $ P.zip (P.reverse codeReversed) (fmap (instructionSize *) [CodeOffset 0 .. ])
let
codeBuilder = mconcat $ fmap (word32LE . getInstruction) code
txt = toLazyByteString $ codeBuilder
<> builderRepeatZero (fromIntegral $ poolOffsetAligned - poolOffset)
<> mconcat (P.reverse poolReversed)
-- resolve symbolTable
let
ff :: (String, Label) -> ElfSymbolXX 'ELFCLASS64
ff (s, r) =
let
steName = s
steBind = STB_Global
steType = STT_NoType
steShNdx = textSecN
steValue = fromIntegral $ findOffset poolOffset r
steSize = 0
in
ElfSymbolXX{..}
symbolTable = ff <$> P.reverse symbolsRefersed
(symbolTableData, stringTableData) <- serializeSymbolTable ELFDATA2LSB (zeroIndexStringItem : symbolTable)
return $ SELFCLASS64 :&: ElfList
[ ElfHeader
{ ehData = ELFDATA2LSB
, ehOSABI = ELFOSABI_SYSV
, ehABIVersion = 0
, ehType = ET_REL
, ehMachine = EM_AARCH64
, ehEntry = 0
, ehFlags = 0
}
, ElfSection
{ esName = ".text"
, esType = SHT_PROGBITS
, esFlags = SHF_EXECINSTR .|. SHF_ALLOC
, esAddr = 0
, esAddrAlign = 8
, esEntSize = 0
, esN = textSecN
, esLink = 0
, esInfo = 0
, esData = ElfSectionData txt
}
, ElfSection
{ esName = ".shstrtab"
, esType = SHT_STRTAB
, esFlags = 0
, esAddr = 0
, esAddrAlign = 1
, esEntSize = 0
, esN = shstrtabSecN
, esLink = 0
, esInfo = 0
, esData = ElfSectionDataStringTable
}
, ElfSection
{ esName = ".symtab"
, esType = SHT_SYMTAB
, esFlags = 0
, esAddr = 0
, esAddrAlign = 8
, esEntSize = symbolTableEntrySize ELFCLASS64
, esN = symtabSecN
, esLink = fromIntegral strtabSecN
, esInfo = 1
, esData = ElfSectionData symbolTableData
}
, ElfSection
{ esName = ".strtab"
, esType = SHT_STRTAB
, esFlags = 0
, esAddr = 0
, esAddrAlign = 1
, esEntSize = 0
, esN = strtabSecN
, esLink = 0
, esInfo = 0
, esData = ElfSectionData stringTableData
}
, ElfSectionTable
]