content: {-# LANGUAGE DeriveAnyClass    #-}
  {-# LANGUAGE DeriveFunctor     #-}
  {-# LANGUAGE DeriveGeneric     #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Kempe.Asm.X86.Type ( X86 (..)
                            , Addr (..)
                            , AbsReg (..)
                            , X86Reg (..)
                            , Label
                            , prettyAsm
                            , prettyDebugAsm
                            ) where
  
  import           Control.DeepSeq         (NFData)
  import qualified Data.ByteString         as BS
  import qualified Data.ByteString.Lazy    as BSL
  import           Data.Copointed
  import           Data.Int                (Int64, Int8)
  import           Data.Text.Encoding      (decodeUtf8)
  import qualified Data.Text.Lazy.Encoding as TL
  import           Data.Word               (Word8)
  import           GHC.Generics            (Generic)
  import           Kempe.Asm.Pretty
  import           Kempe.Asm.Type
  import           Prettyprinter           (Doc, Pretty (pretty), brackets, colon, hardline, (<+>))
  import           Prettyprinter.Ext
  
  type Label = Word
  
  -- | Used to implement 'MovRRLower' pretty-printer
  class As8 reg where
      -- | Return the register for addressing the lower 8 bits given a 64-bit
      -- register
      as8 :: reg -> reg
  
  instance As8 X86Reg where
      as8 R8  = R8b
      as8 R9  = R9b
      as8 R10 = R10b
      as8 R11 = R11b
      as8 R12 = R12b
      as8 R13 = R13b
      as8 R14 = R14b
      as8 R15 = R15b
      as8 Rsi = Sil
      as8 Rdi = Dil
      as8 Rax = AL
      as8 Rdx = DL
      as8 r   = r -- other cases, as8 r8b = r8b, etc.
  
  -- currently just has 64-bit and 8-bit registers
  data X86Reg = R8
              | R9
              | R10
              | R11
              | R12
              | R13
              | R14
              | R15
              | Rdi -- can I use rsi/rdi??
              | Rsi
              -- -- | BH
              -- -- | BL
              | R8b
              | R9b
              | R10b
              | R11b
              | R12b
              | R13b
              | R14b
              | R15b
              | Sil
              | Dil
              | Rsp
              | Rbp
              | Rbx
              -- cl is reserved in this implementation which it really shouldn't be
              -- rax and rdx (and friends) are reserved for unsigned mult.
              | Rcx
              | CH
              | CL
              -- Rax, Rdx and friends are reserved for integer division (and
              -- unsigned multiplication lol)
              | Rax
              | Rdx
              | AH
              | AL
              | DH
              | DL
              deriving (Eq, Ord, Enum, Bounded, Generic, NFData)
  
  instance Pretty X86Reg where
      pretty Rax  = "rax"
      pretty Rcx  = "rcx"
      pretty Rdx  = "rdx"
      pretty Rsp  = "rsp"
      pretty Rbp  = "rbp"
      pretty AH   = "ah"
      pretty AL   = "al"
      pretty CH   = "ch"
      pretty CL   = "cl"
      pretty DH   = "dh"
      pretty DL   = "dl"
      pretty Rbx  = "rbx"
      pretty R8   = "r8"
      pretty R9   = "r9"
      pretty R10  = "r10"
      pretty R11  = "r11"
      pretty R12  = "r12"
      pretty R13  = "r13"
      pretty R14  = "r14"
      pretty R15  = "r15"
      pretty R8b  = "r8b"
      pretty R9b  = "r9b"
      pretty R10b = "r10b"
      pretty R11b = "r11b"
      pretty R12b = "r12b"
      pretty R13b = "r13b"
      pretty R14b = "r14b"
      pretty R15b = "r15b"
      pretty Rsi  = "rsi"
      pretty Rdi  = "rdi"
      pretty Sil  = "sil"
      pretty Dil  = "dil"
  
  data AbsReg = DataPointer
              | AllocReg64 !Int -- TODO: register by size
              | AllocReg8 !Int
              | CArg1
              | CArg2
              | CArg3
              | CArg4
              | CArg5
              | CArg6
              | CRet
              | ShiftExponent
              | QuotRes -- quotient register for idiv, rax
              | RemRes -- remainder register for idiv, rdx
              deriving (Eq, Ord, Generic, NFData)
  
  instance Pretty AbsReg where
      pretty DataPointer    = "datapointer"
      pretty (AllocReg64 i) = "r" <> pretty i
      pretty (AllocReg8 i)  = "HL" <> pretty i
      pretty CRet           = "rax"
      pretty CArg1          = "rdi"
      pretty CArg2          = "rsi"
      pretty CArg3          = "rdx"
      pretty CArg4          = "rcx"
      pretty CArg5          = "r8"
      pretty CArg6          = "r9"
      pretty ShiftExponent  = "cl"
      pretty QuotRes        = "rax"
      pretty RemRes         = "rdx"
  
  -- [ebx+ecx*4h-20h]
  data Addr reg = Reg reg
                | AddrRRPlus reg reg
                | AddrRCPlus reg Int64
                | AddrRCMinus reg Int64
                | AddrRRScale reg reg Int64
                deriving (Generic, NFData)
  
  -- TODO: sanity-check pass to make sure no Reg8's are in e.g. MovRCBool
  
  -- parametric in @reg@; we do register allocation second
  data X86 reg a = PushReg { ann :: a, rSrc :: reg }
                 | PushMem { ann :: a, addr :: Addr reg }
                 | PopMem { ann :: a, addr :: Addr reg }
                 | PopReg { ann :: a, reg :: reg }
                 | PushConst { ann :: a, iSrc :: Int64 }
                 | Jump { ann :: a, label :: Label }
                 | Call { ann :: a, label :: Label }
                 | CallBS { ann :: a, bslLabel :: BSL.ByteString }
                 | Ret { ann :: a }
                 -- intel-ish syntax; destination first
                 | MovRA { ann :: a, rDest :: reg, addrSrc :: Addr reg }
                 | MovAR { ann :: a, addrDest :: Addr reg, rSrc :: reg }
                 | MovABool { ann :: a, addrDest :: Addr reg, boolSrc :: Word8 }
                 | MovRR { ann :: a, rDest :: reg, rSrc :: reg } -- for convenience
                 | MovRRLower { ann :: a, rDest :: reg, rSrc :: reg } -- ^ Doesn't correspond 1-1 to an instruction, writes a 64-bit register to an 8-bit register
                 | MovRC { ann :: a, rDest :: reg, iSrc :: Int64 }
                 | MovRL { ann :: a, rDest :: reg, bsLabel :: BS.ByteString }
                 | MovAC { ann :: a, addrDest :: Addr reg, iSrc :: Int64 }
                 | MovACi8 { ann :: a, addrDest :: Addr reg, i8Src :: Int8 }
                 | MovACTag { ann :: a, addrDest :: Addr reg, tagSrc :: Word8 }
                 | MovRCBool { ann :: a, rDest :: reg, boolSrc :: Word8 }
                 | MovRCi8 { ann :: a, rDest :: reg, i8Src :: Int8 }
                 | MovRCTag { ann :: a, rDest :: reg, tagSrc :: Word8 }
                 | MovRWord { ann :: a, rDest :: reg, wSrc :: Word }
                 | AddRR { ann :: a, rAdd1 :: reg, rAdd2 :: reg }
                 | SubRR { ann :: a, rSub1 :: reg, rSub2 :: reg }
                 | XorRR { ann :: a, rXor1 :: reg, rXor2 :: reg }
                 | ImulRR { ann :: a, rMul1 :: reg, rMul2 :: reg }
                 | AddAC { ann :: a, addrAdd1 :: Addr reg, iAdd2 :: Int64 }
                 | AddRC { ann :: a, rAdd1 :: reg, iAdd2 :: Int64 }
                 | SubRC { ann :: a, rSub1 :: reg, iSub2 :: Int64 }
                 | LShiftLRR { ann :: a, rDest :: reg, rSrc :: reg }
                 | LShiftRRR { ann :: a, rDest :: reg, rSrc :: reg }
                 | AShiftRRR { ann :: a, rDest :: reg, rSrc :: reg }
                 | Label { ann :: a, label :: Label }
                 | BSLabel { ann :: a, bsLabel :: BS.ByteString }
                 | Je { ann :: a, jLabel :: Label }
                 | Jne { ann :: a, jLabel :: Label }
                 | Jg { ann :: a, jLabel :: Label }
                 | Jge { ann :: a, jLabel :: Label }
                 | Jl { ann :: a, jLabel :: Label }
                 | Jle { ann :: a, jLabel :: Label }
                 | CmpAddrReg { ann :: a, addrCmp :: Addr reg, rCmp :: reg }
                 | CmpRegReg { ann :: a, rCmp :: reg, rCmp' :: reg } -- for simplicity
                 | CmpAddrBool { ann :: a, addrCmp :: Addr reg, bCmp :: Word8 }
                 | CmpRegBool { ann :: a, rCmp :: reg, bCmp :: Word8 }
                 | IdivR { ann :: a, rDiv :: reg }
                 | DivR { ann :: a, rDiv :: reg }
                 | Cqo { ann :: a }
                 | AndRR { ann :: a, rDest :: reg, rSrc :: reg }
                 | OrRR { ann :: a, rDest :: reg, rSrc :: reg }
                 | PopcountRR { ann :: a, rDest :: reg, rSrc :: reg }
                 | NegR { ann :: a, rSrc :: reg }
                 | NasmMacro0 { ann :: a, macroName :: BS.ByteString }
                 deriving (Generic, NFData, Functor)
  
  instance Copointed (X86 reg) where
      copoint = ann
  
  instance Pretty reg => Pretty (Addr reg) where
      pretty (Reg r)               = brackets (pretty r)
      pretty (AddrRRPlus r0 r1)    = brackets (pretty r0 <> "+" <> pretty r1)
      pretty (AddrRCPlus r c)      = brackets (pretty r <> "+" <> pretty c)
      pretty (AddrRCMinus r c)     = brackets (pretty r <> "-" <> pretty c)
      pretty (AddrRRScale r0 r1 c) = brackets (pretty r0 <> "+" <> pretty r1 <> "*" <> pretty c)
  
  prettyLive :: (As8 reg, Pretty reg) => X86 reg Liveness -> Doc ann
  prettyLive r = pretty r <+> pretty (ann r)
  
  -- intel syntax
  instance (As8 reg, Pretty reg) => Pretty (X86 reg a) where
      pretty (PushReg _ r)        = i4 ("push" <+> pretty r)
      pretty (PushMem _ a)        = i4 ("push" <+> pretty a)
      pretty (PopMem _ a)         = i4 ("pop qword" <+> pretty a)
      pretty (PopReg _ r)         = i4 ("pop" <+> pretty r)
      pretty (PushConst _ i)      = i4 ("push" <+> pretty i)
      pretty (Jump _ l)           = i4 ("jmp" <+> prettyLabel l)
      pretty (Call _ l)           = i4 ("call" <+> prettyLabel l)
      pretty Ret{}                = i4 "ret"
      pretty (MovRA _ r a)        = i4 ("mov" <+> pretty r <> "," <+> pretty a)
      pretty (MovAR _ a r)        = i4 ("mov" <+> pretty a <> "," <+> pretty r)
      pretty (MovABool _ a b)     = i4 ("mov byte" <+> pretty a <> "," <+> pretty b)
      pretty (MovACi8 _ a i)      = i4 ("mov byte" <+> pretty a <> "," <+> pretty i)
      pretty (MovRCi8 _ r i)      = i4 ("mov byte" <+> pretty r <> "," <+> pretty i)
      pretty (MovRWord _ r w)     = i4 ("mov qword" <+> pretty r <> "," <+> prettyHex w)
      pretty (MovRR _ r0 r1)      = i4 ("mov" <+> pretty r0 <> "," <+> pretty r1)
      pretty (MovRRLower _ r0 r1) = i4 ("mov" <+> pretty r0 <> "," <+> pretty (as8 r1))
      pretty (MovRC _ r i)        = i4 ("mov" <+> pretty r <> "," <+> pretty i)
      pretty (MovAC _ a i)        = i4 ("mov qword" <+> pretty a <> "," <+> pretty i)
      pretty (MovRCBool _ r b)    = i4 ("mov" <+> pretty r <> "," <+> pretty b)
      pretty (MovRL _ r bl)       = i4 ("mov" <+> pretty r <> "," <+> pretty (decodeUtf8 bl))
      pretty (AddRR _ r0 r1)      = i4 ("add" <+> pretty r0 <> "," <+> pretty r1)
      pretty (AddAC _ a c)        = i4 ("add" <+> pretty a <> "," <+> pretty c)
      pretty (SubRR _ r0 r1)      = i4 ("sub" <+> pretty r0 <> "," <> pretty r1)
      pretty (ImulRR _ r0 r1)     = i4 ("imul" <+> pretty r0 <> "," <+> pretty r1)
      pretty (XorRR _ r0 r1)      = i4 ("xor" <+> pretty r0 <> "," <+> pretty r1)
      pretty (AddRC _ r0 c)       = i4 ("add" <+> pretty r0 <> "," <+> pretty c)
      pretty (SubRC _ r0 c)       = i4 ("sub" <+> pretty r0 <> "," <+> pretty c)
      pretty (Label _ l)          = prettyLabel l <> colon
      pretty (BSLabel _ b)        = let pl = pretty (decodeUtf8 b) in "global" <+> pl <> hardline <> pl <> colon
      pretty (Je _ l)             = i4 ("je" <+> prettyLabel l)
      pretty (Jl _ l)             = i4 ("jl" <+> prettyLabel l)
      pretty (CmpAddrReg _ a r)   = i4 ("cmp" <+> pretty a <> "," <+> pretty r)
      pretty (CmpRegReg _ r0 r1)  = i4 ("cmp" <+> pretty r0 <> "," <+> pretty r1)
      pretty (CmpAddrBool _ a b)  = i4 ("cmp byte" <+> pretty a <> "," <+> pretty b)
      pretty (CmpRegBool _ r b)   = i4 ("cmp" <+> pretty r <> "," <+> pretty b)
      pretty (LShiftRRR _ r0 r1)  = i4 ("shr" <+> pretty r0 <> "," <+> pretty r1)
      pretty (LShiftLRR _ r0 r1)  = i4 ("shl" <+> pretty r0 <> "," <+> pretty r1)
      pretty (AShiftRRR _ r0 r1)  = i4 ("sar" <+> pretty r0 <> "," <+> pretty r1)
      pretty (IdivR _ r)          = i4 ("idiv" <+> pretty r)
      pretty (DivR _ r)           = i4 ("div" <+> pretty r)
      pretty Cqo{}                = i4 "cqo"
      pretty (MovACTag _ a t)     = i4 ("mov byte" <+> pretty a <> "," <+> pretty t)
      pretty (AndRR _ r0 r1)      = i4 ("and" <+> pretty r0 <> "," <+> pretty r1)
      pretty (OrRR _ r0 r1)       = i4 ("or" <+> pretty r0 <> "," <+> pretty r1)
      pretty (PopcountRR _ r0 r1) = i4 ("popcnt" <+> pretty r0 <> "," <+> pretty r1)
      pretty (NegR _ r)           = i4 ("neg" <+> pretty r)
      pretty (Jne _ l)            = i4 ("jne" <+> prettyLabel l)
      pretty (Jg _ l)             = i4 ("jg" <+> prettyLabel l)
      pretty (Jge _ l)            = i4 ("jge" <+> prettyLabel l)
      pretty (Jle _ l)            = i4 ("jle" <+> prettyLabel l)
      pretty (MovRCTag _ r b)     = i4 ("mov" <+> pretty r <> "," <+> pretty b)
      pretty (NasmMacro0 _ b)     = i4 (pretty (decodeUtf8 b))
      pretty (CallBS _ b)         = i4 ("call" <+> pretty (TL.decodeUtf8 b))
  
  prettyAsm :: (As8 reg, Pretty reg) => [X86 reg a] -> Doc ann
  prettyAsm = ((prolegomena <#> macros <#> "section .text" <> hardline) <>) . prettyLines . fmap pretty
  
  prettyDebugAsm :: (As8 reg, Pretty reg) => [X86 reg Liveness] -> Doc ann
  prettyDebugAsm = prettyLines . fmap prettyLive
  
  prolegomena :: Doc ann
  prolegomena = "section .bss" <> hardline <> "kempe_data: resb 0x8012" -- 32 kb
  
  macros :: Doc ann
  macros = prettyLines
      [ calleeSave
      , calleeRestore
      , callerSave
      , callerRestore
      ]
  
  -- rbx, rbp, r12-r15 callee-saved (non-volatile)
  -- rest caller-saved (volatile)
  
  -- | Save non-volatile registers
  calleeSave :: Doc ann
  calleeSave =
      "%macro calleesave 0"
      <#> prettyLines (fmap pretty toPush)
      <#> "%endmacro"
      where toPush = PushReg () <$> [Rbx, Rbp, R12, R13, R14, R15]
  
  calleeRestore :: Doc ann
  calleeRestore =
      "%macro calleerestore 0"
      <#> prettyLines (fmap pretty toPop)
      <#> "%endmacro"
      where toPop = PopReg () <$> [R15, R14, R13, R12, Rbp, Rbx]
  
  callerSave :: Doc ann
  callerSave =
      "%macro callersave 0"
      <#> prettyLines (fmap pretty toPush)
      <#> "%endmacro"
      where toPush = PushReg () <$> [Rax, Rcx, Rdx, Rsi, Rdi, R8, R9, R10, R11]
  
  callerRestore :: Doc ann
  callerRestore =
      "%macro callerrestore 0"
      <#> prettyLines (fmap pretty toPop)
      <#> "%endmacro"
      where toPop = PopReg () <$> [R11, R10, R9, R8, Rdi, Rsi, Rdx, Rcx, Rax]
  