content: {-# LANGUAGE DeriveAnyClass    #-}
  {-# LANGUAGE DeriveGeneric     #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  -- | IR loosely based on Appel book.
  module Kempe.IR.Type ( Stmt (..)
                       , Exp (..)
                       , RelBinOp (..)
                       , IntBinOp (..)
                       , BoolBinOp (..)
                       , Label
                       , Temp (..)
                       , WriteSt (..)
                       ) where
  
  import           Control.DeepSeq      (NFData)
  import qualified Data.ByteString      as BS
  import qualified Data.ByteString.Lazy as BSL
  import           Data.Int             (Int64, Int8)
  import           Data.Text.Encoding   (decodeUtf8)
  import           Data.Word            (Word8)
  import           GHC.Generics         (Generic)
  import           Kempe.AST.Size
  import           Prettyprinter        (Doc, Pretty (pretty), braces, brackets, colon, hardline, parens, (<+>))
  import           Prettyprinter.Ext
  
  data WriteSt = WriteSt { wlabels :: [Label]
                         , temps   :: [Int]
                         }
  
  type Label = Word
  
  prettyLabel :: Label -> Doc ann
  prettyLabel l = "kmp" <> pretty l
  
  data Temp = Temp64 !Int
            | Temp8 !Int
            | DataPointer -- RBP on x86 and x19 on aarch64?
            deriving (Eq, Generic, NFData)
  
  instance Pretty Temp where
      pretty (Temp64 i)  = "t_" <> pretty i
      pretty (Temp8 i)   = "t8_" <> pretty i
      pretty DataPointer = "datapointer"
  
  instance Pretty Stmt where
      pretty (Labeled l)           = hardline <> prettyLabel l <> colon
      pretty (Jump l)              = parens ("j" <+> prettyLabel l)
      pretty (CCall ty bs)         = parens ("C" <+> pretty (decodeUtf8 (BSL.toStrict bs)) <+> braces (prettyMonoStackType  ty))
      pretty (KCall l)             = parens ("call" <+> prettyLabel l)
      pretty Ret                   = parens "ret"
      pretty (MovTemp t e)         = parens ("movtemp" <+> pretty t <+> pretty e)
      pretty (MovMem e _ e')       = parens ("movmem" <+> pretty e <+> pretty e') -- TODO: maybe print size?
      pretty (CJump e l l')        = parens ("cjump" <+> pretty e <+> prettyLabel l <+> prettyLabel l')
      pretty (WrapKCall _ ty fn l) = hardline <> "export" <+> pretty (decodeUtf8 fn) <+> braces (prettyMonoStackType ty) <+> prettyLabel l
      pretty (MJump e l)           = parens ("mjump" <+> pretty e <+> prettyLabel l)
  
  instance Pretty Exp where
      pretty (ConstInt i)           = parens ("int" <+> pretty i)
      pretty (ConstInt8 i)          = parens ("int8" <+> pretty i)
      pretty (ConstWord n)          = parens ("word" <+> pretty n)
      pretty (ConstBool False)      = parens "bool false"
      pretty (ConstBool True)       = parens "bool true"
      pretty (Reg t)                = parens ("reg" <+> pretty t)
      pretty (Mem sz e)             = parens ("mem" <+> brackets (pretty sz) <+> pretty e)
      pretty (ExprIntBinOp op e e') = parens (pretty op <+> pretty e <+> pretty e')
      pretty (ExprIntRel op e e')   = parens (pretty op <+> pretty e <+> pretty e')
      pretty (ConstTag b)           = parens ("tag" <+> prettyHex b)
      pretty (BoolBinOp op e e')    = parens (pretty op <+> pretty e <+> pretty e')
      pretty (IntNegIR e)           = parens ("~" <+> pretty e)
      pretty (PopcountIR e)         = parens ("popcount" <+> pretty e)
      pretty (EqByte e e')          = parens ("=b" <+> pretty e <+> pretty e')
  
  data Stmt = Labeled Label
            | Jump Label
            | CJump Exp Label Label -- ^ If the 'Exp' evaluates to @1@, go to the first label, otherwise go to the second (if-then-else). Used to implement ifs.
            | MJump Exp Label
            | CCall MonoStackType BSL.ByteString
            | KCall Label -- ^ a 'KCall' is a jump to a Kempe procedure
            | WrapKCall ABI MonoStackType BS.ByteString Label
            | MovTemp Temp Exp -- ^ Put @e@ in temp
            | MovMem Exp Int64 Exp -- ^ Store @e2@ at address given by @e1@, with sizing information
            | Ret
            deriving (Generic, NFData)
  
  data Exp = ConstInt Int64
           | ConstInt8 Int8
           | ConstTag Word8 -- ^ Used to distinguish constructors of a sum type
           | ConstWord Word
           | ConstBool Bool
           | Reg Temp -- TODO: size?
           | Mem Int64 Exp -- ^ Fetch from address
           | ExprIntBinOp IntBinOp Exp Exp
           | ExprIntRel RelBinOp Exp Exp
           | BoolBinOp BoolBinOp Exp Exp
           | IntNegIR Exp
           | PopcountIR Exp
           | EqByte Exp Exp
           deriving (Eq, Generic, NFData)
             -- TODO: one for data, one for C ABI
  
  data BoolBinOp = BoolAnd
                 | BoolOr
                 | BoolXor
                 deriving (Eq, Generic, NFData)
  
  instance Pretty BoolBinOp where
      pretty BoolAnd = "&"
      pretty BoolOr  = "||"
      pretty BoolXor = "xor"
  
  data RelBinOp = IntEqIR
                | IntNeqIR
                | IntLtIR
                | IntGtIR
                | IntLeqIR
                | IntGeqIR
                deriving (Eq, Generic, NFData)
  
  instance Pretty RelBinOp where
      pretty IntEqIR  = "="
      pretty IntNeqIR = "!="
      pretty IntLtIR  = "<"
      pretty IntGtIR  = ">"
      pretty IntLeqIR = "<="
      pretty IntGeqIR = ">="
  
  data IntBinOp = IntPlusIR
                | IntTimesIR
                | IntDivIR
                | IntMinusIR
                | IntModIR -- rem?
                | IntXorIR
                | WordShiftRIR -- ^ compiles to @shr@ on x86
                | WordShiftLIR
                -- int/word mod are different, see: https://stackoverflow.com/questions/8231882/how-to-implement-the-mod-operator-in-assembly
                | WordModIR
                | WordDivIR
                deriving (Eq, Generic, NFData)
  
  instance Pretty IntBinOp where
      pretty IntPlusIR    = "+"
      pretty IntTimesIR   = "*"
      pretty IntDivIR     = "/"
      pretty IntMinusIR   = "-"
      pretty IntModIR     = "%"
      pretty IntXorIR     = "xor"
      pretty WordShiftRIR = ">>"
      pretty WordShiftLIR = "<<"
      pretty WordModIR    = "%~"
      pretty WordDivIR    = "/~"
  