content: {-# LANGUAGE DeriveAnyClass    #-}
  {-# LANGUAGE DeriveFunctor     #-}
  {-# LANGUAGE DeriveGeneric     #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Kempe.Asm.Arm.Type ( Label
                            , ArmReg (..)
                            , AbsReg (..)
                            , Arm (..)
                            , Cond (..)
                            , Addr (..)
                            , prettyAsm
                            , prettyDebugAsm
                            ) where
  
  import           Control.DeepSeq    (NFData)
  import qualified Data.ByteString    as BS
  import           Data.Copointed
  import           Data.Int           (Int64, Int8)
  import           Data.Text.Encoding (decodeUtf8)
  import           Data.Word          (Word16)
  import           GHC.Generics       (Generic)
  import           Kempe.Asm.Pretty
  import           Kempe.Asm.Type
  import           Prettyprinter      (Doc, Pretty (..), brackets, colon, concatWith, hardline, (<+>))
  import           Prettyprinter.Ext  (prettyHex, prettyLines, (<#>), (<~>))
  
  -- | Sort of silly class that prints the 32-bit equivalent of a register.
  class As32 reg where
      as32b :: reg -> Doc ann
  
  -- r0-r7 result registers
  
  data AbsReg = DataPointer
              | AllocReg !Int
              | CArg0 -- x0
              | CArg1
              | CArg2
              | CArg3
              | CArg4
              | CArg5
              | CArg6
              | CArg7 -- x7
              | LinkReg -- so we can save before/after branch-links
              | StackPtr -- so we can save in translation phase
              deriving (Generic, NFData)
  
  instance Pretty AbsReg where
      pretty DataPointer  = "datapointer"
      pretty (AllocReg i) = "Abs" <> pretty i
      pretty CArg0        = "X0"
      pretty CArg1        = "X1"
      pretty CArg2        = "X2"
      pretty CArg3        = "X3"
      pretty CArg4        = "X4"
      pretty CArg5        = "X5"
      pretty CArg6        = "X6"
      pretty CArg7        = "X7"
      pretty LinkReg      = "X30"
      pretty StackPtr     = "SP"
  
  instance As32 AbsReg where
      as32b = pretty
  
  type Label = Word
  
  data ArmReg = X0
              | X1
              | X2
              | X3
              | X4
              | X5
              | X6
              | X7
              | X8
              | X9
              | X10
              | X11
              | X12
              | X13
              | X14
              | X15
              | X16
              | X17
              | X18
              | X19
              | X20
              | X21
              | X22
              | X23
              | X24
              | X25
              | X26
              | X27
              | X28
              | X29
              | X30 -- ^ This is the link register?
              | SP -- ^ Don't use this
              deriving (Enum, Eq, Ord, Generic, NFData)
  
  instance Pretty ArmReg where
      pretty X0  = "x0"
      pretty X1  = "x1"
      pretty X2  = "x2"
      pretty X3  = "x3"
      pretty X4  = "x4"
      pretty X5  = "x5"
      pretty X6  = "x6"
      pretty X7  = "x7"
      pretty X8  = "x8"
      pretty X9  = "x9"
      pretty X10 = "x10"
      pretty X11 = "x11"
      pretty X12 = "x12"
      pretty X13 = "x13"
      pretty X14 = "x14"
      pretty X15 = "x15"
      pretty X16 = "x16"
      pretty X17 = "x17"
      pretty X18 = "x18"
      pretty X19 = "x19"
      pretty X20 = "x20"
      pretty X21 = "x21"
      pretty X22 = "x22"
      pretty X23 = "x23"
      pretty X24 = "x24"
      pretty X25 = "x25"
      pretty X26 = "x26"
      pretty X27 = "x27"
      pretty X28 = "x28"
      pretty X29 = "x29"
      pretty X30 = "x30"
      pretty SP  = "sp"
  
  instance As32 ArmReg where
      as32b X0  = "w0"
      as32b X1  = "w1"
      as32b X2  = "w2"
      as32b X3  = "w3"
      as32b X4  = "w4"
      as32b X5  = "w5"
      as32b X6  = "w6"
      as32b X7  = "w7"
      as32b X8  = "w8"
      as32b X9  = "w9"
      as32b X10 = "w10"
      as32b X11 = "w11"
      as32b X12 = "w12"
      as32b X13 = "w13"
      as32b X14 = "w14"
      as32b X15 = "w15"
      as32b X16 = "w16"
      as32b X17 = "w17"
      as32b X18 = "w18"
      as32b X19 = "w19"
      as32b X20 = "w20"
      as32b X21 = "w21"
      as32b X22 = "w22"
      as32b X23 = "w23"
      as32b X24 = "w24"
      as32b X25 = "w25"
      as32b X26 = "w26"
      as32b X27 = "w27"
      as32b X28 = "w28"
      as32b X29 = "w29"
      as32b X30 = "w30"
      as32b SP  = error "Internal error: as32b sp should not happen!!"
  
  data Addr reg = Reg reg
                | AddRRPlus reg reg
                | AddRCPlus reg Int64
                deriving (Eq, Generic, NFData)
  
  instance (Pretty reg) => Pretty (Addr reg) where
      pretty (Reg reg)         = brackets (pretty reg)
      pretty (AddRRPlus r0 r1) = brackets (pretty r0 <~> pretty r1)
      pretty (AddRCPlus r i)   = brackets (pretty r <~> prettyInt i)
  
  -- | See: https://developer.arm.com/documentation/dui0068/b/arm-instruction-reference/conditional-execution?lang=en
  data Cond = Eq | Neq | UnsignedLeq | UnsignedGeq | UnsignedLt
            | Geq | Lt | Gt | Leq deriving (Generic, NFData)
  
  instance Pretty Cond where
      pretty Eq          = "EQ"
      pretty Neq         = "NE"
      pretty UnsignedLeq = "LS"
      pretty Geq         = "GE"
      pretty Lt          = "LT"
      pretty Gt          = "GT"
      pretty Leq         = "LE"
      pretty UnsignedLt  = "LO"
  
  -- | For reference: https://static.docs.arm.com/100898/0100/the_a64_Instruction_set_100898_0100.pdf
  --
  -- https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions?lang=en
  data Arm reg a = Branch { ann :: a, label :: Label } -- like jump
                 | BranchLink { ann :: a, label :: Label } -- like @call@
                 | BranchCond { ann :: a, label :: Label, cond :: Cond }
                 | BranchZero { ann :: a, condReg :: reg, label :: Label }
                 | BranchNonzero { ann :: a, condReg :: reg, label :: Label }
                 | AddRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg }
                 | AddRC { ann :: a, res :: reg, inp1 :: reg, int :: Int64 }
                 | SubRC { ann :: a, res :: reg, inp1 :: reg, int :: Int64 }
                 | SubRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg }
                 | MulRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg }
                 | MulSubRRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg, inp3 :: reg }
                 | MovRC { ann :: a, dest :: reg, iSrc :: Int64 } -- TODO: change this to a Word16
                 | SignedDivRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg }
                 | UnsignedDivRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg }
                 | MovRWord { ann :: a, dest :: reg, wSrc :: Word16 }
                 | MovRK { ann :: a, dest :: reg, wSrc :: Word16, lShift :: Int8 }
                 | MovRR { ann :: a, dest :: reg, src :: reg }
                 | AndRR { ann :: a, dest :: reg, inp1 :: reg, inp2 :: reg }
                 | OrRR { ann :: a, dest :: reg, inp1 :: reg, inp2 :: reg }
                 | XorRR { ann :: a, dest :: reg, inp1 :: reg, inp2 :: reg }
                 | Load { ann :: a, dest :: reg, addrSrc :: Addr reg }
                 | LoadByte { ann :: a, dest :: reg, addrSrc :: Addr reg }
                 | LoadLabel { ann :: a, dest :: reg, srcLabel :: BS.ByteString }
                 | Store { ann :: a, src :: reg, addrDest :: Addr reg }
                 | StoreByte { ann :: a, src :: reg, addrDest :: Addr reg } -- ^ @strb@ in Aarch64 assembly, "store byte"
                 | CmpRR { ann :: a, inp1 :: reg, inp2 :: reg }
                 | CmpRC { ann :: a, inp1 :: reg, int :: Int64 }
                 | CSet { ann :: a, dest :: reg, cond :: Cond }
                 | Ret { ann :: a }
                 | Label { ann :: a, label :: Label }
                 | BSLabel { ann :: a, bsLabel :: BS.ByteString }
                 | LShiftLRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg } -- LShift - logical shift
                 | LShiftRRR { ann :: a, res :: reg, inp1 :: reg, inp2 :: reg }
                 | GnuMacro { ann :: a, macroName :: BS.ByteString }
                 | Neg { ann :: a, dest :: reg, src :: reg }
                 deriving (Functor, Generic, NFData)
  
  -- | Don't call this on a negative number!
  prettyUInt :: Integral a => a -> Doc b
  prettyUInt i = "#" <> prettyHex i
  
  prettyInt :: (Pretty a) => a -> Doc b
  prettyInt = ("#" <>) . pretty
  
  instance (Pretty reg, As32 reg) => Pretty (Arm reg a) where
      pretty (Branch _ l)              = i4 ("b" <+> prettyLabel l)
      pretty (BranchLink _ l)          = i4 ("bl" <+> prettyLabel l)
      pretty (BranchCond _ l c)        = i4 ("b." <> pretty c <+> prettyLabel l)
      pretty (BranchZero _ r l)        = i4 ("cbz" <+> pretty r <~> prettyLabel l)
      pretty (BranchNonzero _ r l)     = i4 ("cbnz" <+> pretty r <~> prettyLabel l)
      pretty Ret{}                     = i4 "ret"
      pretty (BSLabel _ b)             = let pl = pretty (decodeUtf8 b) in ".globl" <+> pl <> hardline <> pl <> colon
      pretty (MovRWord _ r c)          = i4 ("mov" <+> pretty r <~> prettyUInt c)
      pretty (MovRK _ r c l)           = i4 ("movk" <+> pretty r <~> prettyUInt c <~> "lsl" <+> pretty l)
      pretty (LShiftLRR _ r r0 r1)     = i4 ("lsl" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (LShiftRRR _ r r0 r1)     = i4 ("lsr" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (AddRR _ r r0 r1)         = i4 ("add" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (SubRR _ r r0 r1)         = i4 ("sub" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (MulRR _ r r0 r1)         = i4 ("mul" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (MulSubRRR _ r r0 r1 r2)  = i4 ("msub" <+> pretty r <~> pretty r0 <~> pretty r1 <~> pretty r2)
      pretty (SignedDivRR _ r r0 r1)   = i4 ("sdiv" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (UnsignedDivRR _ r r0 r1) = i4 ("udiv" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (Load _ r a)              = i4 ("ldr" <+> pretty r <~> pretty a)
      pretty (LoadByte _ r a)          = i4 ("ldrb" <+> as32b r <~> pretty a)
      pretty (LoadLabel _ r l)         = i4 ("ldr" <+> pretty r <~> "=" <> pretty (decodeUtf8 l))
      pretty (Store _ r a)             = i4 ("str" <+> pretty r <~> pretty a)
      pretty (StoreByte _ r a)         = i4 ("strb" <+> as32b r <~> pretty a)
      pretty (MovRR _ r0 r1)           = i4 ("mov" <+> pretty r0 <~> pretty r1)
      pretty (AndRR _ r r0 r1)         = i4 ("and" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (OrRR _ r r0 r1)          = i4 ("orr" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (XorRR _ r r0 r1)         = i4 ("eor" <+> pretty r <~> pretty r0 <~> pretty r1)
      pretty (CSet _ r c)              = i4 ("cset" <+> pretty r <~> pretty c)
      pretty (MovRC _ r i)             = i4 ("mov" <+> pretty r <~> prettyInt i)
      pretty (CmpRR _ r0 r1)           = i4 ("cmp" <+> pretty r0 <~> pretty r1)
      pretty (Label _ l)               = prettyLabel l <> colon
      pretty (GnuMacro _ b)            = i4 (pretty (decodeUtf8 b))
      pretty (AddRC _ r r0 i)          = i4 ("add" <+> pretty r <~> pretty r0 <~> "#" <> pretty i)
      pretty (SubRC _ r r0 i)          = i4 ("sub" <+> pretty r <~> pretty r0 <~> "#" <> pretty i)
      pretty (CmpRC _ r0 i)            = i4 ("cmp" <+> pretty r0 <~> "#" <> pretty i)
      pretty (Neg _ r0 r1)             = i4 ("neg" <+> pretty r0 <~> pretty r1)
  
  instance Copointed (Arm reg) where
      copoint = ann
  
  prettyAsm :: (Pretty reg, As32 reg) => [Arm reg a] -> Doc ann
  prettyAsm = (<> hardline) . ((prolegomena <#> macros <#> ".text" <> hardline) <>) . prettyLines . fmap pretty
  
  -- http://www.mathcs.emory.edu/~cheung/Courses/255/Syl-ARM/7-ARM/array-define.html
  prolegomena :: Doc ann
  prolegomena = ".p2align 3" <#> ".data" <#> "kempe_data: .skip 32768" -- 32kb
  -- see https://stackoverflow.com/questions/65748033/how-to-fix-ld-warning-arm64-function-not-4-byte-aligned-warning
  -- since mine are byte-aligned
  
  macros :: Doc ann
  macros = prettyLines
      [ calleeSave
      , calleeRestore
      , callerSave
      , callerRestore
      ]
  
  -- see:
  -- https://community.arm.com/developer/ip-products/processors/b/processors-ip-blog/posts/using-the-stack-in-aarch64-implementing-push-and-pop
  
  calleeSave :: Doc ann
  calleeSave =
      ".macro calleesave"
      <#> i4 "sub sp, sp, #(8 * 10)" -- allocate space on stack
      <#> prettyLines (fmap pretty stores)
      <#> ".endm"
      where toPush = [X19 .. X28]
            stores = zipWith (\r o -> Store () r (AddRCPlus SP (8*o))) toPush [0..]
  
  calleeRestore :: Doc ann
  calleeRestore =
      ".macro calleerestore"
      <#> prettyLines (fmap pretty loads)
      <#> i4 "add sp, sp, #(8 * 10)" -- free stack space
      <#> ".endm"
      where toPop = [X19 .. X28]
            loads = zipWith (\r o -> Load () r (AddRCPlus SP (8*o))) toPop [0..]
  
  callerSave :: Doc ann
  callerSave =
      ".macro callersave"
      <#> i4 "sub sp, sp, #(8 * 8)"
      <#> prettyLines (fmap pretty stores)
      <#> ".endm"
      where toPush = X30 : [X9 .. X15]
            stores = zipWith (\r o -> Store () r (AddRCPlus SP (8*o))) toPush [0..]
  
  callerRestore :: Doc ann
  callerRestore =
      ".macro callerrestore"
      <#> prettyLines (fmap pretty loads)
      <#> i4 "add sp, sp, #(8 * 8)"
      <#> ".endm"
      where toPop = X30 : [X9 .. X15]
            loads = zipWith (\r o -> Load () r (AddRCPlus SP (8*o))) toPop [0..]
  
  prettyLive :: (As32 reg, Pretty reg) => Arm reg Liveness -> Doc ann
  prettyLive r = pretty r <+> pretty (ann r)
  
  prettyDebugAsm :: (As32 reg, Pretty reg) => [Arm reg Liveness] -> Doc ann
  prettyDebugAsm = concatWith (<#>) . fmap prettyLive
  