hasm-revised-addressmode.hs

{-#LANGUAGE GADTs, DataKinds, KindSignatures #-}

import Data.Word
import Data.Int
import Numeric
import Data.Bits
import Data.Binary.Put (putWord32le, putWord16le, runPut)
import Data.ByteString.Lazy (unpack)
import qualified Data.ByteString as B

-- | Binary string to word. Use like fB "101" ==> 5
fB :: String -> Word8
fB = fromIntegral . fst . head . readInt 2 (`elem` "01") convertion
    where convertion '0' = 0
          convertion '1' = 1

data Extent = With | Without

data R32 ::  Extent -> Extent -> * where
    EAX :: R32 e e
    ECX :: R32 e e
    EDX :: R32 e e
    EBX :: R32 e e
    ESI :: R32 e e
    EDI :: R32 e e
    ESP :: R32 With e
    EBP :: R32 e With


regToIndex :: R32 a b -> Word8
regToIndex EAX = fB "000"
regToIndex ECX = fB "001"
regToIndex EDX = fB "010"
regToIndex EBX = fB "011"
regToIndex ESP = fB "100"
regToIndex EBP = fB "101"
regToIndex ESI = fB "110"
regToIndex EDI = fB "111"

data Scale = Scale1 | Scale2 | Scale4 | Scale8

data AdressMode a b =
    EADirect (R32 a b)
  | EAIndirect (R32 Without Without)
  | EAIndirectDisp8 (R32 Without b) Int8
  | EAIndirectDisp32 (R32 Without b) Int32
  | Disp32 Int32
  | SIB Scale (R32 Without a) (R32 a Without)
  | SIBDisp8 Scale (R32 Without a) (R32 a Without) Int8
  | SIBDisp32 Scale (R32 Without a) (R32 a Without) Int32
  | SIDisp32 Scale (R32 Without a) Int32
  | SIEBPDisp8 Scale (R32 Without a) Int8
  | SIEBPDisp32 Scale (R32 Without a) Int32

threePart :: Word8 -> Word8 -> Word8 -> Word8  
threePart mod reg rm = mod `shiftL` 6 .|. (reg `shiftL` 3) .|. rm

getModRM :: AdressMode a b -> R32 c d -> Word8
getModRM (EAIndirect r) b                 = threePart (fB "00") (regToIndex b) (regToIndex r)
getModRM (EAIndirectDisp8 r _) b          = threePart (fB "01") (regToIndex b) (regToIndex r)
getModRM (EAIndirectDisp32 r _) b         = threePart (fB "10") (regToIndex b) (regToIndex r)
getModRM (EADirect r) b                   = threePart (fB "11") (regToIndex b) (regToIndex r)

getModRM (Disp32 _) b                     = threePart (fB "00") (regToIndex b) (fB "101")
getModRM (SIB _ _ _) b                    = threePart (fB "00") (regToIndex b) (fB "100")
getModRM (SIBDisp8 _ _ _ _) b             = threePart (fB "01") (regToIndex b) (fB "100")
getModRM (SIBDisp32 _ _ _ _) b            = threePart (fB "10") (regToIndex b) (fB "100")

getModRM (SIDisp32 _ _ _) b               = threePart (fB "00") (regToIndex b) (fB "100")
getModRM (SIEBPDisp8 _ _ _) b             = threePart (fB "01") (regToIndex b) (fB "100")
getModRM (SIEBPDisp32 _ _ _) b            = threePart (fB "10") (regToIndex b) (fB "100")

mkSIBbyte :: Scale -> R32 a b -> R32 c d -> Maybe Word8
mkSIBbyte scale index base = Just $ threePart (scaleToWord8 scale) (regToIndex index) (regToIndex base)
    where
        scaleToWord8 Scale1 = fB "00"
        scaleToWord8 Scale2 = fB "01"
        scaleToWord8 Scale4 = fB "10"
        scaleToWord8 Scale8 = fB "11"

getSIB :: AdressMode a b -> Maybe Word8
getSIB (SIB scale index base)          = mkSIBbyte scale index base
getSIB (SIBDisp8 scale index base _)   = mkSIBbyte scale index base
getSIB (SIBDisp32 scale index base _)  = mkSIBbyte scale index base
getSIB (SIDisp32 scale index _)        = mkSIBbyte scale index EBP
getSIB (SIEBPDisp8 scale index _)      = mkSIBbyte scale index EBP
getSIB (SIEBPDisp32 scale index _)     = mkSIBbyte scale index EBP
getSIB _                               = Nothing

int32ToBytes :: Int32 -> [Word8]
int32ToBytes = unpack . runPut . putWord32le . fromIntegral

getDisplacement :: AdressMode a b -> [Word8]
getDisplacement (EAIndirectDisp8 _ disp)          = [fromIntegral disp]
getDisplacement (EAIndirectDisp32 _ disp)         = int32ToBytes disp
getDisplacement (Disp32 disp)                     = int32ToBytes disp
getDisplacement (SIBDisp8 _ _ _ disp)             = [fromIntegral disp]
getDisplacement (SIBDisp32 _ _ _ disp)            = int32ToBytes disp
getDisplacement (SIDisp32 _ _ disp)               = int32ToBytes disp
getDisplacement (SIEBPDisp8 _ _ disp)             = [fromIntegral disp]
getDisplacement (SIEBPDisp32 _ _ disp)            = int32ToBytes disp

mkModRM_SIB_Displacement :: AdressMode a b -> R32 c d -> [Word8]
mkModRM_SIB_Displacement mode reg = 
    case getSIB mode of
        Just sib -> getModRM mode reg : sib : getDisplacement mode
        Nothing -> getModRM mode reg : getDisplacement mode