{-  CAO Compiler
    Copyright (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>. -}

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE BangPatterns               #-}

{- | 
Module      :  $Header$
Description :  CAO AST and constructors.
Copyright   :  (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho
License     :  GPL
   
Maintainer  :  Paulo Silva <paufil@di.uminho.pt>
Stability   :  experimental
Portability :  non-portable (<reason>)
  
This module contains the definition of data types which represent the
abstract syntax of the CAO language.

The abstract syntax is parametric on the identifier (variables),
allowing for storing different information accordingly with the phase.

Each data type includees its respective pretty printer instance.

-}

module Language.CAO.Syntax where

import Data.Foldable ( Foldable )

import Data.Traversable ( Traversable )

import Language.CAO.Common.Literal
import Language.CAO.Common.Operator
import Language.CAO.Common.Outputable
import Language.CAO.Common.Polynomial
import Language.CAO.Common.SrcLoc

import Language.CAO.Type

--------------------------------------------------------------------------------
-- * Programs
--------------------------------------------------------------------------------

-- | A CAO program is list of definitions with source code location annotations.
-- The order of the definition is relevant since the type checking expects 
-- that required symbols are in the context. 
data Prog id = Prog 
    { unDefs :: [LDef id]
    , initP  :: Maybe (Fun id) }
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (Prog id) where
    ppr = pprProg

pprProg :: PP id => Prog id -> CDoc
pprProg (Prog defs ini) = vsep (map ppr defs)
    $+$ maybe empty ppr ini

-- * Definitions and declarations

-- | A CAO 'Located' definition 'Def'
--
type LDef id = Located (Def id)

-- | A CAO Definition
--
data Def id
    -- | Global variable definition
    = VarDef   (VarDecl id)
    -- | Global constant definition
    | ConstDef (ConstDecl id)
    -- | Function/procedure definition 
    | FunDef   (Fun id)
    -- | Type definition
    | TyDef    (TyDef id)
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (Def id) where
    ppr = pprDef

pprDef :: PP id => Def id -> CDoc
pprDef (VarDef   vd) = ppr vd
pprDef (ConstDef cd) = ppr cd
pprDef (FunDef   fd) = ppr fd
pprDef (TyDef    td) = ppr td

-- | There are three types of CAO variable declarations: 
--
-- 1) Simple variable declaration. 
--
--    This is a single variable declaration with an optional initialization
--     of the form:
--
--  * Without initialization. @VarD var tydecl Nothing@ where @var@ is an
--     identifier with location information and @tydecl@ is its type.
--
--     @def v1 : int;@
--
--  * With initialization. @VarD var tydecl (Just expr)@ now includes an
--     initialization expression, @expr@.
--
--     @def v2 : int := 3;@
--
-- 2) Multiple variable declarations with the same type.
--     @MultiD vars tydecl@ where @vars@ is a (non-empty) list of identifiers 
--     with location information and @tydecl@ is their type.
--
--     @def v1, v2, ..., vn : int;@
--
-- 3) Container variable declarations with initialization.
--     @CondD var tydecl init@ where @var@ is an identifier with location
--     information, @tydecl@ is its type and @init@ is a (non-empty) list of 
--     initialization expressions. Moreover, the length of the list must
--     correspond to the container size in the type declaration.
--
--     @def v1 : vector [3] of int := { 1, 2, 3 };@
--
data VarDecl id
    -- | Simple variable declaration
    = VarD (Located id)       -- Variable declared
           (TyDecl id)        -- Type of the variables
           (Maybe (TLExpr id)) -- Initializer
    -- | Multiple variable declaration
    | MultiD [Located id]     -- Variables declared
             (TyDecl id)      -- Type of the variables
    -- | Container variable declaration
    | ContD (Located id)      -- Variable
            (TyDecl id)       -- Type
            [TLExpr id]        -- Container elems
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (VarDecl id) where
    ppr = pprVarDecl

pprVarDecl :: PP id => VarDecl id -> CDoc
pprVarDecl (VarD v ty Nothing)
    =   text "def" <+> ppr v <+> colon <+> ppr ty <> semi
pprVarDecl (VarD v ty (Just e))
    =   text "def" <+> ppr v <+> colon <+> ppr ty 
    <+> text ":="  <+> ppr e <> semi
pprVarDecl (MultiD vs ty)
    =   text "def" <+> pprElems vs <+> colon <+> ppr ty <> semi
pprVarDecl (ContD v ty es)
    =   text "def" <+> ppr v <+> colon <+> ppr ty
    <+> text ":="  <+> braces (pprElems es) <> semi

-- | Symbolic CAO constant declarations:
--
-- 1) Single constant declarations. @ConstD var tydecl ann@ where @var@ is
-- an identifier with location, @tydecl@ is its type and @ann@ is an 
-- additional annotation about the constant as defined in 'ConstAnn':
--
-- * With no additional information:
--
--   @def const c : int;@
--
-- * With a constant initialization:
--
--   @def const c : int := 3;@
--
--   @def const c : int := 3 * a;@
--
-- * With a condition (invariant):
--
--   @def const c : int { 0 < c };@
--
-- 2) Multiple constant declarations:
-- 
-- * Without a condition: @MultiConstD vars tydecl Nothing@ where @vars@ is a
--    (non-empty) list of identifiers with location information and @tydecl@ is
--    their type:
--
--   @def const c1, c2, ..., cn : int;@
--
-- * With a condition @MultiConstD vars typdecl (Just cond)@ now includes a
--    condition expression @cond@ that is an invariant:
--
--   @def const c1, c2, ..., cn : int { c1 < c2 && c2 < cn };@
--
data ConstDecl id
    -- | Single constant declaration
    = ConstD      (Located id)       -- Constant declared
                  (TyDecl id)        -- Type of the constant
                  (ConstAnn id)      -- Additional information about constant
    -- | Multiple constant declaration
    | MultiConstD [Located id]       -- Constants declared
                  (TyDecl id)        -- Type of the variables
                  (Maybe (LExpr id)) -- Optional condition
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (ConstDecl id) where
    ppr = pprConstDecl

pprConstDecl :: PP id => ConstDecl id -> CDoc
pprConstDecl (ConstD c ty an)
    =   text "def" <+> text "const" <+> ppr c <+> colon <+> ppr ty
    <+> ppr an <> semi
pprConstDecl (MultiConstD cs ty Nothing)
    =   text "def" <+> text "const" <+> pprElems cs <+> colon <+> ppr ty <> semi
pprConstDecl (MultiConstD cs ty (Just cond))
    =   text "def" <+> text "const" <+> pprElems cs <+> colon <+> ppr ty 
    <+> pprCond cond <> semi

pprCond :: PP id => LExpr id -> CDoc
pprCond c = lbrace <+> ppr c <+> rbrace 

-- | Additional information about symbolic constants.
data ConstAnn id
    -- | No additional information required.
    = None                  
    -- | Definition of the symbolic constant using other constants.
    | ConstInit (LExpr id)
    -- | Condition (invariant) that the symbolic constant must obey.
    | ConstCond (LExpr id)
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (ConstAnn id) where
    ppr = pprConstAnn

pprConstAnn :: PP id => ConstAnn id -> CDoc
pprConstAnn None             = empty
pprConstAnn (ConstInit e)    = text ":=" <+> ppr e
pprConstAnn (ConstCond cond) = ppr cond

-- | A CAO Funtion definition
--
data Fun id
    = Fun { funId      :: Located id  -- ^ Function identifier with location
                                      --   information.
          , funArgs    :: [Arg id]    -- ^ List of function parameters (can be
                                      -- empty).
          , returnType :: [TyDecl id] -- ^ List of the types of returned values.
                                      -- This can be empty in the case of
                                      -- procedures.
          , funBody    :: [LStmt id]  -- ^ Function body. The list of statements
                                      -- must be non-empty.
          }
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (Fun id) where
    ppr = pprFun

pprFun :: PP id => Fun id -> CDoc
pprFun (Fun fn args rt stmts)
    = flip pprBlock stmts $ text "def" <+> ppr fn <> parens (pprElems args)
                                       <+> colon <+> pprRetTys_ rt 

pprRetTys_ :: PP id => [TyDecl id] -> CDoc
pprRetTys_ []  = text "void"
pprRetTys_ tys = pprElems tys

-- | Function arguments/parameters:
--
-- 1) Regular arguments. @Arg var tydecl@ where @var@ is an identifier with
--     location and @tydecl@ is its type:
--
--     @(..., a : int, ...)@
--
-- 2) Symbolic (constant) arguments:
-- 
-- * Without a condition. @ArgConst var tydecl Nothing@ where @var@ is an
--    identifier and @tydecl@ is its type:
--
--    @(..., const a : int, ...)@
-- 
-- * With a condition. @ArgConst var tydecl (Just cond)@ now includes a
--    condition expression @cond@ which is an invariant in the body of the
--    function:
--
--    @)(..., const a : int { 0 < a }, ...)@
--
data Arg id 
    -- | Regular argument declaration
    = Arg      (Located id) (TyDecl id)
    -- | Symbolic argument declaration
    | ArgConst (Located id) (TyDecl id) (Maybe (LExpr id))
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (Arg id) where
    ppr = pprArg

pprArg :: PP id => Arg id -> CDoc
pprArg (Arg n ty)
    = ppr n <+> colon <+> ppr ty
pprArg (ArgConst n ty Nothing) 
    = text "const" <+> ppr n <+> colon <+> ppr ty
pprArg (ArgConst n ty (Just c)) 
    = text "const" <+> ppr n <+> colon <+> ppr ty <+> pprCond c

-- | A CAO type synonym declaration or struct definition.
--
-- 1) A type synonym, @TySynDef var tydecl@ where @var@ is an identifier
--     with location information of the type alias and @tydecl@ is the
--     definition of the type.
--
--  @typedef V8 := vector[8] of int;@
--
-- 2) A structure declaration, 
--     @StructDecl var [(p1, tydecl1), ... (pn, tydecln)]@
--     where @var@ is an indentifier with location information of the structure,
--     @pi@ with @i@ in @1,...,n@, are identifier with location information of
--     the struture projections and
--     @tydecli@ with @i@ in @1,...,n@ are the respective projection types.
--
--     @typedef S := struct [ def f1 : int; 
--                      def f2 : bool; ];@
--
data TyDef id
    -- | Type synonym
    = TySynDef   (Located id)
                 (TyDecl id)
    -- | Struct declaration
    | StructDecl (Located id)
                 [(Located id, TyDecl id)] -- Struct fields
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (TyDef id) where
    ppr = pprTyDef

pprTyDef :: PP id => TyDef id -> CDoc
pprTyDef (TySynDef tn ty)
  =  text "typedef" <+> ppr tn <+> text ":=" <+> ppr ty <> semi
pprTyDef (StructDecl sn flds)
  =  text "typedef" <+> ppr sn <+> text ":=" <+> text "struct"
  <> pprFlds_ flds <> semi

pprFlds_ :: PP id => [(Located id, TyDecl id)] -> CDoc
pprFlds_ flds = brackets $ sep $ map pprFld_ flds

pprFld_ :: PP id => (Located id, TyDecl id) -> CDoc
pprFld_ (fld, ty) = text "def" <+> ppr fld <+> colon <+> ppr ty <> semi

-- | A 'Located' 'TyDecl'
--
type LTyDecl id = Located (TyDecl id)

-- | A CAO type declaration. It is the type described in, for example:
-- @
-- def x : int;
-- def y : mod [ 2 ];
-- ...
-- @
data TyDecl id
    = IntD                  -- ^ Arbitrary precision integers
    | RIntD                 -- ^ Machine (register) precision integers
    | BoolD                 -- ^ Booleans
    | BitsD Sign (LExpr id) -- ^ Bits strings with /sign/ and /length/
    | ModD (Mod id)         -- ^ Modular type
    | VectorD (LExpr id)    
              (TyDecl id)   -- ^ Vectors with /size/ and /type/ of elements
    | MatrixD (LExpr id)    
              (LExpr id)  
              (TyDecl id)   -- ^ Matrices with /row/ and /column sizes/ and 
                            -- /type/ of elements
    | TySynD (Located id)   -- ^ Type alias
    deriving (Functor, Foldable, Traversable, Show, Eq, Read)

instance PP id => PP (TyDecl id) where
    ppr = pprTyDecl

pprTyDecl :: PP id => TyDecl id -> CDoc
pprTyDecl IntD
    = text "int"
pprTyDecl RIntD
    = text "register" <+> text "int"
pprTyDecl BoolD
    = text "bool"
pprTyDecl (BitsD snm e)
    = ppr snm <+> text "bits" <> brackets (ppr e)
pprTyDecl (ModD md)
    = text "mod" <> brackets (ppr md)
pprTyDecl (VectorD ln ty)
    = text "vector" <> brackets (ppr ln) <+> text "of" <+> ppr ty
pprTyDecl (MatrixD rows cols ty)
    = text "matrix" <> brackets (pprElems [rows,cols]) <+> text "of" <+> ppr ty
pprTyDecl (TySynD sid)
    = ppr sid

-- | A CAO Modulus
data Mod id
    -- | Numeric modulus: @mod [ <int expr> ]@
    = ModNum (LExpr id)
    -- | Modular extention: @mod [ <type><ind> / <pol> ]@
    | ModPol (TyDecl id) id (Pol id)
    deriving (Functor, Foldable, Traversable, Show, Eq, Read)

instance PP id => PP (Mod id) where
    ppr = pprMod

pprMod :: PP id => Mod id -> CDoc
pprMod (ModNum n)
    = ppr n
pprMod (ModPol ty n pol)
    = ppr ty <> char '<' <> ppr n <> char '>' <+> char '/' <+> ppr pol

-- * Expressions

-- | CAO Located Expressions
--
type LExpr id = Located (Expr id)

-- | CAO Located Expressions with type annotations
type TLExpr id = Located (TExpr id)

-- | CAO Expressions
data Expr id
    -- | Variable
    = Var id
    -- | Literal
    | Lit (Literal id)
    -- | Function call
    | FunCall (Located id) [TLExpr id]
    -- | Struct projection
    | StructProj (TLExpr id) id
    -- | Unary operation
    | UnaryOp UOp (TLExpr id)
    -- | Binary operation
    | BinaryOp (BinOp id) (TLExpr id) (TLExpr id)
    -- | Container access
    | Access (TLExpr id) (APat id)
    -- | Cast
    | Cast Bool [LTyDecl id] (TLExpr id)
    deriving (Functor, Foldable, Traversable, Show, Eq, Read)

-- | Unary operations
data UOp
    -- | Symmetric expression (@-@)
    = Sym
    -- | Boolean negation (@!@)
    | Not
    -- | Bitwise negation (@~@)
    | BNot
    deriving (Eq, Show, Read)

instance PP UOp where
    ppr = pprUOp

pprUOp :: UOp -> CDoc
pprUOp Sym  = char '-'
pprUOp Not  = char '!'
pprUOp BNot = char '~'

-- | Binary operations
data BinOp id
    -- | Arithmetic binary operation
    = ArithOp AOp
    -- | Boolean binary operation
    | BoolOp BOp
    -- | Bitwise operation
    | BitOp BWOp
    -- | Shift/rotate operation. This case is separated because the
    -- second argument is an index.
    | BitsSROp SROp 
    -- | Comparison
    | CmpOp (Type id) COp
    -- | Concat operation (@\@@)
    | Concat
    deriving (Eq, Show, Read, Functor, Foldable, Traversable)

instance PP (BinOp id) where
    ppr = pprBinOp

pprBinOp :: BinOp id -> CDoc
pprBinOp (ArithOp  op) = ppr op
pprBinOp (BoolOp   op) = ppr op
pprBinOp (BitOp    op) = ppr op
pprBinOp (BitsSROp op) = ppr op
pprBinOp (CmpOp _  op) = ppr op
pprBinOp Concat        = char '@' 

instance PP id => PP (Expr id) where
    ppr = pprExpr

pprExpr :: PP id => Expr id -> CDoc
pprExpr (Var v)
    = ppr v
pprExpr (Lit l)
    = ppr l
pprExpr (FunCall fn es)
    = ppr fn <> parens (pprElems es)
pprExpr ctx@(StructProj (L _ e) fi)
    = pprParens_ e ctx <> char '.' <> ppr fi
pprExpr ctx@(UnaryOp op (L _ e))
    = ppr op <> pprParens_ e ctx
pprExpr ctx@(BinaryOp op (L _ l) (L _ r))
    = pprParensL_ l ctx <+> ppr op <+> pprParensR_ r ctx
pprExpr ctx@(Access (L _ l) p)
    = pprParens_ l ctx <> ppr p
pprExpr ctx@(Cast True td (L _ e))
    = parens (ppr td) <> pprParens_ e ctx
pprExpr (Cast False _ e)
    = ppr e

-- | Arithmetic Operations
--
data AOp 
    = Plus  -- ^ Sum (@+@)
    | Minus -- ^ Subtraction (@-@)
    | Times -- ^ Multiplication (@*@)
    | Power -- ^ Exponentiation (@**@)
    | Div   -- ^ Whole division (@/@)
    | ModOp -- ^ Remainer of whole division (@%@)
    deriving (Eq, Show, Read)

instance PP AOp where
    ppr = pprAOp

pprAOp :: AOp -> CDoc
pprAOp Plus  = char '+'
pprAOp Minus = char '-'
pprAOp Times = char '*'
pprAOp Power = text "**"
pprAOp Div   = char '/'
pprAOp ModOp = char '%'

-- | Comparision operations
--
data COp 
    = Eq    -- ^ Equality (@==@)
    | Neq   -- ^ Not equal (@!=@)
    | Lt    -- ^ Less than (@<@)
    | Leq   -- ^ Less than or equal (@<=@)
    | Gt    -- ^ Greater than (@>@)
    | Geq   -- ^ Greater than or equal (@>=@)
    deriving (Eq, Show, Read)

instance PP COp where
    ppr = pprCOp

pprCOp :: COp -> CDoc
pprCOp Eq  = text "==" 
pprCOp Neq = text "!="
pprCOp Lt  = char '<'
pprCOp Leq = text "<="
pprCOp Gt  = char '>'
pprCOp Geq = text ">="

-- | Boolean operations
--
data BOp 
    = And   -- ^ Boolean conjunction (@&&@)
    | Or    -- ^ Boolean disjunction (@||@)
    | Xor   -- ^ Boolean exclusive disjunction (\xor\) (@^^@)
    deriving (Eq, Show, Read)

instance PP BOp where
    ppr = pprBOp

pprBOp :: BOp -> CDoc
pprBOp And = text "&&" 
pprBOp Or  = text "||"
pprBOp Xor = text "^^"

-- | Bit string operations
--
data BWOp
    = BWOr  -- ^ Bitwise disjunction (@|@)
    | BWAnd -- ^ Bitwise conjunction (@&@)
    | BWXor -- ^ Bitwise exclusive disjunction (@^@)
    deriving (Eq, Show, Read)

instance PP BWOp where
    ppr = pprBWOp

pprBWOp :: BWOp -> CDoc
pprBWOp BWAnd = char '&'
pprBWOp BWOr  = char '|'
pprBWOp BWXor = char '^'

-- | Shifts/Rotate operations
--
data SROp 
    = SUp   -- ^ Shift up (left) (@<<@)
    | SDown -- ^ Shift down (right) (@>>@)
    | RUp   -- ^ Rotate up  (@<|@)
    | RDown -- ^ Rotate down (@|>@)
    deriving (Eq, Show, Read)

instance PP SROp where
    ppr = pprSROp

pprSROp :: SROp -> CDoc
pprSROp SUp   = text "<<"
pprSROp SDown = text ">>"
pprSROp RUp   = text "<|"
pprSROp RDown = text "|>"

-- | Vector and matrice accesses
--
data APat id
    -- | Vector access/range
    = VectP (RowAPat id)
    -- | Matrix access/range (row, column)
    | MatP  (RowAPat id) (ColAPat id)
    deriving (Functor, Foldable, Traversable, Show, Eq, Read)

instance PP id => PP (APat id) where
    ppr = pprAPat

pprAPat :: PP id => APat id -> CDoc
pprAPat (VectP r)   = brackets $ ppr r
pprAPat (MatP  r c) = brackets (ppr r <> comma <> ppr c)

-- | Accesses and ranges
--
data RowAPat id
    -- | Element access using an integer expression as index
    = CElem  (TLExpr id)
    -- | Range access using two integer expressions as limits of the range:
    -- @CRange from to@.
    | CRange (TLExpr id) 
             (TLExpr id)
    deriving (Functor, Foldable, Traversable, Show, Eq, Read)

instance PP id => PP (RowAPat id) where
    ppr = pprRowAPat

pprRowAPat :: PP id => RowAPat id -> CDoc
pprRowAPat (CElem e)    = ppr e
pprRowAPat (CRange f t) = ppr f <> text ".." <> ppr t

-- | Column access pattern.
--
type ColAPat = RowAPat

-- * Statements

-- | AST Annotations. Annotations have a no semantics and 
-- are not pretty-printed nor translated. They are used solely for internal
-- processing purposes.
data Annot 
    = EndIndex
    | EndAux
    | EndConsts

-- | CAO Located Statement
--
type LStmt id = Located (Stmt id)

-- | CAO Statements
--
data Stmt id
    -- | Variable declaration
    = VDecl  (VarDecl id)
    -- | Constant declaration
    | CDecl  (ConstDecl id)
    -- | Assignment
    | Assign [LVal id] [TLExpr id]
    -- | Function call
    | FCallS id [TLExpr id]
    -- | Return statement
    | Ret [TLExpr id]
    -- | Conditional statement (If)
    | Ite (TLExpr id)
          [LStmt id]
          (Maybe [LStmt id])
    -- | Sequence statement (Seq)
    | Seq (SeqIter id) [LStmt id]

    -- | Iterative statement (While)
    | While (TLExpr id) [LStmt id]
    -- | No operation (not syntactic)
    | Nop Annot
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (Stmt id) where
    ppr = pprStmt

pprStmt :: PP id => Stmt id -> CDoc
pprStmt (VDecl vd)
    = ppr vd
pprStmt (CDecl cd)
    = ppr cd
pprStmt (Assign lvs es)
    = pprElems lvs <+> text ":=" <+> pprElems es <> semi
pprStmt (FCallS fn es)
    = ppr fn <> parens (pprElems es) <> semi
pprStmt (Ret es)
    = text "return" <+> pprElems es <> semi
pprStmt (Ite c i me)
    = pprBlock (text "if" <> parens (ppr c)) i <> pprElse me
    where 
    pprElse Nothing  = empty
    pprElse (Just e) = pprBlock (text "else") e
pprStmt (Seq s b)
    = pprBlock (ppr s) b
pprStmt (While e st)
    = pprBlock (text "while" <> parens (ppr e)) st
pprStmt (Nop _) 
    = empty

-- | Sequence iterations
data SeqIter id
    = SeqIter { seqVar    :: id               -- ^ Bound variable identifier
              , seqStart  :: LExpr id         -- ^ Start index
              , seqEnd    :: LExpr id         -- ^ End index
              , seqBy     :: Maybe (LExpr id) -- ^ Optional increment
              , seqIdx    :: SeqRange         -- ^ Expansion annotation
              }
    deriving (Functor, Foldable, Traversable)

-- | Sequence annotations
data SeqRange
    = SimpleRng [Integer]   -- ^ Simple sequence indexes
    | NestedRng [[Integer]] -- ^ Nested sequence indexes
    deriving (Eq, Show)

instance PP id => PP (SeqIter id) where
    ppr = pprSeqIter

pprSeqIter :: PP id => SeqIter id -> CDoc
pprSeqIter (SeqIter v s e mb _)
    =   text "seq" <+> ppr v <+> text ":="
    <+> ppr s <+> text "to" <+> ppr e <+> pprBy mb
    where 
    pprBy Nothing  = empty
    pprBy (Just b) = text "by" <+> ppr b

-- | CAO left value
data LVal id
    -- | Simple left value @x := ...@
    = LVVar (Located id)
    -- | Struct left value @s.fi := ...@
    | LVStruct (LVal id) id
    -- | Containers @v[0] := ... m[1,2] := ...@
    | LVCont (Type id) (LVal id) (APat id)
    deriving (Functor, Foldable, Traversable)

instance PP id => PP (LVal id) where
    ppr = pprLVal

pprLVal :: PP id => LVal id -> CDoc
pprLVal (LVVar v)        = ppr v
pprLVal (LVStruct sn fi) = ppr sn <> char '.' <> ppr fi
pprLVal (LVCont _ c p)   = ppr c <> ppr p

-- * Type annotations

data TExpr id = TyE (Type id) (Expr id)
    deriving (Show, Read, Eq, Foldable, Traversable, Functor)

instance PP id => PP (TExpr id) where
    ppr (TyE _ e) = ppr e

{-# INLINE annTyE #-}
annTyE :: Type id -> Expr id -> TExpr id
annTyE t e = TyE t e

{-# INLINE annL #-}
annL :: Type id -> LExpr id -> TLExpr id
annL t = fmap (annTyE t)

{-# INLINE unTyp #-}
unTyp :: TExpr id -> Expr id
unTyp (TyE _ e) = e

{-# INLINE unTypL #-}
unTypL :: TLExpr id -> LExpr id
unTypL = fmap unTyp

--------------------------------------------------------------------------------
-- UTILS
--------------------------------------------------------------------------------

pprBlock :: PP id => CDoc -> [LStmt id] -> CDoc
pprBlock header stmts
    = header <+> lbrace $+$ nest 2 (vcat $ map ppr stmts) $+$ rbrace

instance Operator (TExpr id) where
    isSimple (TyE _ e) = isSimple e
    assoc (TyE _ e) = assoc e
    fixity (TyE _ e) = fixity e
    prec (TyE _ e) = prec e

instance Operator (Expr id) where
    isSimple (Lit _)                  = True
    isSimple (Var _)                  = True 
    isSimple (Cast False _ (L _ e))   = isSimple e
    isSimple _                        = False

    assoc (Cast False _ (L _ e))      = assoc e
    assoc _                           = ALeft

    fixity (UnaryOp _ _)              = Prefix
    fixity (Cast True _ _)            = Prefix
    fixity (FunCall _ _)              = Postfix
    fixity (StructProj _ _)           = Postfix
    fixity (Access {})                = Postfix
    fixity (Cast False _ (L _ e))     = fixity e
    fixity _                          = Infix

    prec (Var _)                      = 200
    prec (Lit _)                      = 200
    prec (FunCall _ _)                = 190
    prec (StructProj _ _)             = 190
    prec (UnaryOp _ _)                = 180
    prec (BinaryOp (ArithOp op) _ _)  = prec op
    prec (BinaryOp (BoolOp op) _ _)   = prec op
    prec (BinaryOp (BitOp op) _ _)    = prec op
    prec (BinaryOp (BitsSROp op) _ _) = prec op
    prec (BinaryOp (CmpOp _ op) _ _)  = prec op
    prec (BinaryOp Concat _ _)        = 150
    prec (Access {})                  = 190
    prec (Cast True _ _)              = 170
    prec (Cast False _ (L _ e))       = prec e

instance Operator AOp where
    isSimple _ = False
    assoc _    = ALeft
    fixity _   = Infix
    prec Plus  = 140
    prec Minus = 140
    prec Times = 150
    prec Div   = 150
    prec ModOp = 150
    prec Power = 160

instance Operator COp where
    isSimple _ = False
    assoc _    = ALeft
    fixity _   = Infix
    prec Eq    = 110
    prec Neq   = 110
    prec Lt    = 120
    prec Leq   = 120
    prec Gt    = 120
    prec Geq   = 120

instance Operator BOp where
    isSimple _ = False
    assoc _    = ALeft
    fixity _   = Infix
    prec And   = 60
    prec Or    = 40
    prec Xor   = 50

instance Operator BWOp where
    isSimple _ = False
    assoc _    = ALeft
    fixity _   = Infix
    prec BWOr  = 70
    prec BWAnd = 90
    prec BWXor = 80

instance Operator SROp where
    isSimple _ = False
    assoc _    = ALeft
    fixity _   = Infix
    prec SUp   = 130
    prec SDown = 130
    prec RUp   = 130
    prec RDown = 130