content: module LLVM.Codegen.IR
    ( IR(..)
    , Terminator(..)
    , ComparisonType(..)
    , CallingConvention(..)
    , TailCallAttribute(..)
    , SynchronizationScope(..)
    , MemoryOrdering(..)
    , Alignment
    , Flag(..)
    , NUW
    , NSW
    , Exact
    , Inbounds
    , Volatile
    , renderIR
    ) where
  
  import Prelude hiding (EQ)
  import LLVM.Codegen.Name
  import LLVM.Codegen.Operand
  import LLVM.Codegen.Type
  import LLVM.Codegen.Flag
  import LLVM.Pretty
  import Data.Word
  import Data.List.NonEmpty (NonEmpty(..))
  import qualified Data.List.NonEmpty as NE
  
  
  data NUW
  data NSW
  data Exact
  data Inbounds
  data Volatile
  
  type Alignment = Word32
  
  data SynchronizationScope
    = SingleThread
    | System
    deriving Show
  
  data MemoryOrdering
    = Unordered
    | Monotonic
    | Acquire
    | Release
    | AcquireRelease
    | SequentiallyConsistent
    deriving Show
  
  type Atomicity = (SynchronizationScope, MemoryOrdering)
  
  data ComparisonType
    = EQ
    | NE
    | UGT
    | UGE
    | ULT
    | ULE
    | SGT
    | SGE
    | SLT
    | SLE
    deriving (Eq, Show)
  
  data TailCallAttribute
    = Tail
    | MustTail
    | NoTail
    deriving Show
  
  data CallingConvention
    = C
    | Fast
    -- TODO add others as needed
    deriving Show
  
  data IR
    = Add !(Flag NUW) !(Flag NSW) !Operand !Operand
    | Mul !(Flag NUW) !(Flag NSW) !Operand !Operand
    | Sub !(Flag NUW) !(Flag NSW) !Operand !Operand
    | Udiv !(Flag Exact) !Operand !Operand
    | And !Operand !Operand
    | Or !Operand !Operand
    | Trunc !Operand !Type
    | Zext !Operand !Type
    | Bitcast !Operand !Type
    | ICmp !ComparisonType !Operand !Operand
    | PtrToInt !Operand !Type
    | Alloca !Type !(Maybe Operand) !Int
    | GetElementPtr !(Flag Inbounds) !Operand ![Operand]
    | Load !(Flag Volatile) !Operand !(Maybe Atomicity) !Alignment
    | Store !(Flag Volatile) !Operand !Operand !(Maybe Atomicity) !Alignment
    | Phi !(NonEmpty (Operand, Name))
    | Call !(Maybe TailCallAttribute) !CallingConvention !Operand ![Operand]  -- TODO support param attributes
    -- Terminators
    | Ret !(Maybe Operand)
    | Br !Name
    | CondBr !Operand !Name !Name
    | Switch !Operand !Name ![(Operand, Name)]
    | Select !Operand !Operand !Operand
    deriving Show
  
  newtype Terminator
    = Terminator IR
    deriving Show
  
  renderTCA :: Renderer TailCallAttribute
  renderTCA buf = \case
    Tail -> buf |># "tail"#
    NoTail -> buf |># "notail"#
    MustTail -> buf |># "musttail"#
  {-# INLINABLE renderTCA #-}
  
  renderCC :: Renderer CallingConvention
  renderCC buf = \case
    C -> buf |># "ccc"#
    Fast -> buf |># "fastcc"#
  {-# INLINABLE renderCC #-}
  
  renderComparisonType :: Renderer ComparisonType
  renderComparisonType buf = \case
    EQ  -> buf |># "eq"#
    NE  -> buf |># "ne"#
    UGT -> buf |># "ugt"#
    UGE -> buf |># "uge"#
    ULT -> buf |># "ult"#
    ULE -> buf |># "ule"#
    SGT -> buf |># "sgt"#
    SGE -> buf |># "sge"#
    SLT -> buf |># "slt"#
    SLE -> buf |># "sle"#
  {-# INLINABLE renderComparisonType #-}
  
  renderMemoryOrdering :: Renderer MemoryOrdering
  renderMemoryOrdering buf = \case
    Unordered              -> buf |># "unordered"#
    Monotonic              -> buf |># "monotonic"#
    Acquire                -> buf |># "acquire"#
    Release                -> buf |># "release"#
    AcquireRelease         -> buf |># "acq_rel"#
    SequentiallyConsistent -> buf |># "seq_cst"#
  {-# INLINABLE renderMemoryOrdering #-}
  
  renderSyncScope :: Renderer SynchronizationScope
  renderSyncScope buf = \case
    SingleThread ->
      buf |># "syncscope(\"singlethread\")"#
    System ->
      buf
  {-# INLINABLE renderSyncScope #-}
  
  renderIR :: Renderer IR
  renderIR buf = \case
    Add nuw nsw a b ->
      renderArithBinOp buf "add "# nuw nsw a b
    Mul nuw nsw a b ->
      renderArithBinOp buf "mul "# nuw nsw a b
    Sub nuw nsw a b ->
      renderArithBinOp buf "sub "# nuw nsw a b
    Udiv exact a b ->
      ((((optional exact (buf |># "udiv "#) (|># "exact "#) `renderType` typeOf a)
      |>. ' ') `renderOperand` a) |># ", "#) `renderOperand` b
    And a b ->
      ((((buf |># "and "#) `renderType` typeOf a) |>. ' ') `renderOperand` a |># ", "#) `renderOperand` b
    Or a b ->
      ((((buf |># "or "#) `renderType` typeOf a) |>. ' ') `renderOperand` a |># ", "#) `renderOperand` b
    ICmp cmp a b ->
      (((((buf |># "icmp "#) `renderComparisonType` cmp |>. ' ') `renderType` typeOf a) |>. ' ') `renderOperand` a |># ", "#) `renderOperand` b
    Trunc val to ->
      renderConvertOp buf "trunc "# val to
    Zext val to ->
      renderConvertOp buf "zext "# val to
    Bitcast val to ->
      renderConvertOp buf "bitcast "# val to
    PtrToInt val to ->
      renderConvertOp buf "ptrtoint "# val to
    Alloca ty mNumElems alignment ->
      renderMaybe
        (renderMaybe ((buf |># "alloca "#) `renderType` ty) mNumElems
          (\buf' count -> ((buf' |># ", "#) `renderType` typeOf count |>. ' ') `renderOperand` count))
        (if alignment == 0 then Nothing else Just alignment)
        (\buf' align -> buf' |># ", align "# |>$ align)
    GetElementPtr inbounds pointer indices ->
      case typeOf pointer of
        ty@(PointerType innerTy) ->
          commas (((optional inbounds (buf |># "getelementptr "#) (|># "inbounds "#) `renderType` innerTy |># ", "#) `renderType` ty |>. ' ')
          `renderOperand` pointer |># ", "#) indices prettyIndex
        _ ->
          buf |> error "Operand given to `getelementptr` that is not a pointer!"
      where
        prettyIndex :: Buffer %1 -> Operand -> Buffer
        prettyIndex buf' i = (renderType buf' (typeOf i) |>. ' ') `renderOperand` i
    Load volatile addr atomicity alignment ->
      case atomicity of
        Nothing ->
          withAlignment alignment
            ((((optional volatile (buf |># "load "#) (|># "volatile "#)
              `renderType` resultTy) |># ", "#) `renderType` ptrTy |>. ' ') `renderOperand` addr)
        Just (syncScope, memoryOrdering) ->
          withAlignment alignment
            ((((((optional volatile (buf |># "load atomic "#) (|># "volatile "#)
              `renderType` resultTy) |># ", "#) `renderType` ptrTy |>. ' ') `renderOperand` addr |>. ' ')
              `renderSyncScope` syncScope |>. ' ') `renderMemoryOrdering` memoryOrdering)
      where
        ptrTy = typeOf addr
        resultTy = case ptrTy of
          PointerType ty -> ty
          _ -> error "Malformed AST, expected pointer type."
    Store volatile addr value atomicity alignment ->
      case atomicity of
        Nothing ->
          withAlignment alignment
            ((((optional volatile (buf |># "store "#) (|># "volatile "#) `renderType` ty |>. ' ') `renderOperand` value |># ", "#)
              `renderType` ptrTy |>. ' ') `renderOperand` addr)
        Just (syncScope, memoryOrdering) ->
          withAlignment alignment
            ((((((optional volatile (buf |># "store atomic "#) (|># "volatile "#) `renderType` ty |>. ' ') `renderOperand` value |># ", "#)
              `renderType` ptrTy |>. ' ') `renderOperand` addr |>. ' ') `renderSyncScope` syncScope |>. ' ') `renderMemoryOrdering` memoryOrdering)
      where
        ty = typeOf value
        ptrTy = PointerType ty
    Phi cases@((val, _) :| _) ->
      commas ((buf |># "phi "#) `renderType` typeOf val |>. ' ') (NE.toList cases) renderPhiCase
      where
        renderPhiCase :: Renderer (Operand, Name)
        renderPhiCase buf' (value, name) =
          brackets buf' (\buf'' -> (renderOperand buf'' value |># ", %"#) `renderName` name)
    Call tcAttr cc fn args ->
      (((renderMaybe buf tcAttr (\buf' tca -> renderTCA buf' tca |>. ' ')
        |># "call "#) `renderCC` cc |>. ' ') `renderType` resultType |>. ' ')
        `renderOperand` fn `renderArgs` args
      where
        resultType = case typeOf fn of
          PointerType (FunctionType retTy _) -> retTy
          FunctionType retTy _ -> retTy
          _ -> error "Malformed AST, expected function type."
        renderArgs :: Renderer [Operand]
        renderArgs buf' args' = tupled buf' args' renderArg
        renderArg :: Renderer Operand
        renderArg buf' arg =
          (renderType buf' (typeOf arg) |>. ' ') `renderOperand` arg
    Ret term -> case term of
      Nothing ->
        buf |># "ret void"#
      Just operand ->
        ((buf |># "ret "#) `renderType` typeOf operand |>. ' ') `renderOperand` operand
    Br blockName ->
      (buf |># "br label %"#) `renderName` blockName
    CondBr cond trueLabel falseLabel ->
      (((buf |># "br i1 "#) `renderOperand` cond
        |># ", label %"#) `renderName` trueLabel
        |># ", label %"#) `renderName` falseLabel
    Switch val defaultLabel cases ->
      brackets ((((buf |># "switch "#) `renderType` typeOf val |>. ' ') `renderOperand` val |># ", label %"#) `renderName` defaultLabel |>. ' ')
        (\buf' -> hsep buf' cases renderCase)
      where
        renderCase :: Renderer (Operand, Name)
        renderCase buf' (caseVal, label) =
          ((renderType buf' (typeOf caseVal) |>. ' ') `renderOperand` caseVal |># ", label %"#) `renderName` label
    Select c t f ->
      ((((((buf |># "select "#) `renderType` typeOf c |>. ' ') `renderOperand` c |># ", "#)
        `renderType` typeOf t |>. ' ') `renderOperand` t |># ", "#)
        `renderType` typeOf f |>. ' ') `renderOperand` f
    where
      withAlignment :: Word32 -> Buffer %1 -> Buffer
      withAlignment alignment buf' =
        if alignment == 0
          then buf'
          else buf' |># ", align "# |>$ alignment
  {-# INLINABLE renderIR #-}
  
  renderArithBinOp :: Buffer %1 -> Addr# -> Flag NUW -> Flag NSW -> Operand -> Operand -> Buffer
  renderArithBinOp buf opName nuw nsw a b =
    (((optional nsw (optional nuw
      (buf |># opName) (|># "nuw "#)) (|># "nsw "#) `renderType` typeOf a) |>. ' ') `renderOperand` a |># ", "#) `renderOperand` b
  {-# INLINABLE renderArithBinOp #-}
  
  renderConvertOp :: Buffer %1 -> Addr# -> Operand -> Type -> Buffer
  renderConvertOp buf opName val to =
    ((((buf |># opName) `renderType` typeOf val) |>. ' ') `renderOperand` val |># " to "#) `renderType` to
  {-# INLINABLE renderConvertOp #-}
  