{-  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 FlexibleContexts #-}
{-# LANGUAGE PatternGuards    #-}

{- |
Module      :  $Header$
Description :  CAO type checker.
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

This module implements the typechecker of the CAO language. This follows the paper ...
-}

module Language.CAO.Typechecker ( tcCaoAST ) where

import Control.Applicative ( (<$>) )
import Control.Monad

import Data.List ( genericLength ) 
import qualified Data.Set as Set

import Language.CAO.Common.Error
import Language.CAO.Common.Literal
import Language.CAO.Common.Monad 
import Language.CAO.Common.Polynomial hiding ((.*.), (.+.))
import Language.CAO.Common.SrcLoc
import Language.CAO.Common.State
import Language.CAO.Common.Utils
import Language.CAO.Common.Var

import Language.CAO.Index

import Language.CAO.Syntax
import Language.CAO.Syntax.Utils

import Language.CAO.Type
import Language.CAO.Type.Utils

import Language.CAO.Typechecker.Check
import Language.CAO.Typechecker.Constraint
import Language.CAO.Typechecker.Expr
import Language.CAO.Typechecker.Heap
import Language.CAO.Typechecker.Index
import Language.CAO.Typechecker.PostProcessor
import Language.CAO.Typechecker.Solver

import Main.Flags (RunMode(..))
--------------------------------------------------------------------------------
--Typechecker
--------------------------------------------------------------------------------

{-# INLINE tcError' #-}
tcError' :: CaoMonad m => ErrorCode Var -> m a
tcError' a = tcError a

--Typechecking Prog ------------------------------------------------------------

mapWithLoc :: CaoMonad m => (a -> m b) -> Located a -> m (Located b)
mapWithLoc f (L loc e) = setSrcLoc loc >> liftM (L loc) (f e)

tcCaoAST :: CaoMonad m => RunMode -> Prog Name -> m (Prog Var, Heap)
tcCaoAST m ast = withTcST $ do
    setMode m
    ast' <- tcProg ast
    h <- getHeap
    return (ast', h)

-- Prog ------------------------------------------------------------------------

tcProg :: CaoMonad m => Prog Name -> m (Prog Var)
tcProg (Prog defs _) = liftM2 Prog (mapM (mapWithLoc tcDef) defs) (return Nothing)

-- Definition ------------------------------------------------------------------

tcDef :: CaoMonad m => Def Name -> m (Def Var)
tcDef (VarDef v)   = VarDef . fst            <$> tcVarDecl Global v
tcDef (ConstDef c) = ensureDepMode $ ConstDef <$> tcConstDecl Global c
tcDef (FunDef f)   = FunDef                  <$> tcFunc f
tcDef (TyDef t)    = TyDef                   <$> tcTypeDef t

-- Func ------------------------------------------------------------------------

tcFunc :: CaoMonad m => Fun Name -> m (Fun Var)
tcFunc (Fun (L loc n) args rtype body) = keepGScope $ do
    setSrcLoc loc
    (args', ats)    <- mapAndUnzipM tcArg args
    (rtype', rt)    <- tcTypeDecls rtype
    (body', st, cc) <- tcStmts rt body
    checkFuncReturn st rt
    let tct = FuncSig ats rt cc
    n' <- checkDecl Global tct n
    return (Fun (L loc n') args' rtype' body')

-- Arg -------------------------------------------------------------------------

tcArg :: CaoMonad m => Arg Name -> m (Arg Var, Type Var)

tcArg (Arg (L loc v) td) = do
    setSrcLoc loc
    (td', tct) <- tcTypeDecl td
    v' <- checkDecl Local tct v
    return (Arg (L loc v') td', tct)

tcArg (ArgConst (L loc c) td cond) = ensureDepMode $ do
    setSrcLoc loc
    (td', tct) <- tcIndexDecl td
    c' <- checkConstDecl Local c tct Nothing
    (cond', i') <- withCond cond
    return (ArgConst (L loc c') td' cond', Index c' i' tct)
    where
    withCond Nothing = return (Nothing, Nothing)
    withCond (Just cnd) = do
        (cnd', i) <- tcICond cnd
        addHypothesis [i]
        return (Just cnd', Just i)

--TypeDef ----------------------------------------------------------------------

tcTypeDef :: CaoMonad m => TyDef Name -> m (TyDef Var)
tcTypeDef (TySynDef (L loc n) d) = do
    setSrcLoc loc
    (st', t) <- tcTypeDecl d
    n' <- checkTySyn n t
    return (TySynDef (L loc n') st')
tcTypeDef (StructDecl ln@(L loc _) sfs) = do
    setSrcLoc loc
    (tds, tys) <- mapAndUnzipM tcStructFldDecl sfs
    checkStructDecl ln (map fst sfs) tds tys 
    where
    tcStructFldDecl :: CaoMonad m => (Located Name, TyDecl Name) -> m (TyDecl Var, Type Var)
    tcStructFldDecl (L lc _, td) = setSrcLoc lc >> tcTypeDecl td

--TypeDecl ---------------------------------------------------------------------

tcTypeDecls :: CaoMonad m => [TyDecl Name] -> m ([TyDecl Var], Type Var)
tcTypeDecls = liftM (mapSnd toTuple) . mapAndUnzipM tcTypeDecl 

tcTypeDeclsL :: CaoMonad m => [LTyDecl Name] -> m ([LTyDecl Var], Type Var)
tcTypeDeclsL = liftM (mapSnd toTuple) . mapAndUnzipM (\ (L loc d) -> setSrcLoc loc >> liftM (mapFst (L loc)) (tcTypeDecl d))

tcTypeDecl :: CaoMonad m => TyDecl Name -> m (TyDecl Var, Type Var)
tcTypeDecl IntD =
    return (IntD, Int)
tcTypeDecl RIntD =
    return (RIntD, RInt)
tcTypeDecl BoolD =
    return (BoolD, Bool)
tcTypeDecl (BitsD sign es) = do
    (es', i) <- tcIExpr RInt es
    checkBitsSize i
    return (BitsD sign es', Bits sign i)
tcTypeDecl (ModD (ModNum eb)) = do
    (eb', b) <- tcIExpr Int eb
    checkModBase b
    let tct = Mod Nothing Nothing $ Pol [Mon (CoefI b) EZero]
    return (ModD (ModNum eb'), tct)
tcTypeDecl (ModD (ModPol td ti poly)) = do
    (td', tct) <- tcTypeDecl td
    checkMod tct
    tcWarn (PolExtensionWarn poly)
    when (isModInt tct) $
        tcWarn $ BaseExtensionWarn (getPoly tct)
    (ti', poly', mpt) <- checkPolynomial tct ti poly
    checkPolyLit poly mpt
    return (ModD (ModPol td' ti' poly'), mpt)
tcTypeDecl (VectorD ei td) = do
    (ei', i) <- tcIExpr RInt ei
    checkVectorSize i
    (td', tct) <- tcTypeDecl td
    return (VectorD ei' td', Vector i tct)
tcTypeDecl (MatrixD er ec td) = do
    (er', r) <- tcIExpr RInt er
    (ec', c) <- tcIExpr RInt ec
    checkMatrixSize r c
    (td', tct) <- tcTypeDecl td
    checkAlgebraic tct
    let tct' = Matrix r c tct
    return (MatrixD er' ec' td', tct')
tcTypeDecl (TySynD (L loc n)) = do
    setSrcLoc loc
    v <- checkScopeType n
    return (TySynD (L loc v), synType $ varType v)

-- Statement -------------------------------------------------------------------

-- XXX: The use of maximumClass in this way is not very efficient
tcStmts :: CaoMonad m => Type Var -> [LStmt Name] -> m ([LStmt Var], Type Var, Class Var)
tcStmts rt stmts = do
    (s, t, c) <- fold3M (tcLStmt rt) (:) goTy (\ x  y -> maximumClass [x, y]) ([], Bullet, Pure) stmts
    s' <- cleanDeadCode s
    return (s', t, c)
    where 
    goTy :: Type Var -> Type Var -> Type Var
    goTy Bullet t = t
    goTy t      _ = t
    

cleanDeadCode :: CaoMonad m => [LStmt id] -> m [LStmt id]
cleanDeadCode [] = return []
cleanDeadCode (e@(L l (Ret _)) : r) = do
    unless (null r) $ withSrcLoc l (tcWarn (DeadCodeReturn :: WarningCode Var))  
    return [e]
cleanDeadCode (e : r) = liftM (e :) $ cleanDeadCode r

tcLStmt :: CaoMonad m => Type Var -> LStmt Name -> m (LStmt Var, Type Var, Class Var)
tcLStmt rt (L loc stmt) = withSrcLoc loc $ do
    (stmt', ty, c) <- tcStmt stmt rt
    return (L loc stmt', ty, c)

tcStmt :: CaoMonad m => Stmt Name -> Type Var -> m (Stmt Var, Type Var, Class Var)
tcStmt (VDecl vd) _ = do
    (vd', c) <- tcVarDecl Local vd
    return (VDecl vd', Bullet, c)
tcStmt (CDecl cd) _ = ensureDepMode $ do
    cd' <- tcConstDecl Local cd
    return (CDecl cd', Bullet, Pure)
tcStmt (Assign lvs es@[L fcl (TyE _ (FunCall (L nl n) exps))]) _ = do
    v <- checkScopeFunc n
    if isProc (varType v)
        then tcProc lvs fcl nl v exps
        else tcAssign lvs es
tcStmt (Assign lvs es) _ = 
    tcAssign lvs es
tcStmt (FCallS fid exps) _ = do
    fid' <- checkScopeProc fid
    let FuncSig pts frt (Proc wvars) = varType fid'
    unless (isNil frt) $ tcError (BadUseException fid' ProcScope)
    (exps', pts', cc, _) <- checkArgs exps pts
    let wvars' = Set.toList $ wVars cc  `Set.union` Set.fromList wvars
        fid'' = setType (FuncSig pts' Bullet (Proc wvars)) fid'
    return (FCallS fid'' exps', Bullet, Proc wvars')
tcStmt (Ret exps) rt = do
    (exps'', ccs) <- checkMultipleAssign exps (fromTuple rt)
    return (Ret exps'', rt, maximumClass ccs)
tcStmt (Ite cond istms Nothing) rt = do
    (cond', c1) <- tcExpTyp cond Bool
    (istms', _, c2) <- keepScope $ tcStmts rt istms
    return (Ite cond' istms' Nothing, Bullet, maximumClass [c1, c2])
tcStmt (Ite cond istms (Just estms)) rt = do
    (cond', cb) <- tcExpTyp cond Bool
    (istms', ist, cc1) <- keepScope $ tcStmts rt istms
    (estms', est, cc2) <- keepScope $ tcStmts rt estms
    return ( Ite cond' istms' (Just estms')
           , if not ((isNil ist) || (isNil est)) then rt else Bullet
           , maximumClass [cb, cc1, cc2]
           )
tcStmt (While cond wstms) rt = do
    (cond', cb) <- tcExpTyp cond Bool
    (wstms', _, cc) <- keepScope $ tcStmts rt wstms
    return (While cond' wstms', Bullet, maximum [cb, cc])
tcStmt (Seq (SeqIter ivar estart eend eby _) sstms) rt = do
    (estart', start) <- tcIExpr RInt estart
    (eend', end)     <- tcIExpr RInt eend
    (eby', by)       <- tcBy eby
    keepScope $ do
        ivar' <- checkConstDecl Local ivar RInt Nothing
        let i = IInd $ ivar'
            (start', end') = if by > 0 then (start, end) else (end, start)
        addHypothesis [start' .<=. i, i .<=. end']
        (sstms', sst, c) <- tcStmts rt sstms
        return (Seq (SeqIter ivar' estart' eend' eby' (SimpleRng [])) sstms', sst, c)
tcStmt (Nop a) _ = return (Nop a, Bullet, Pure)

tcBy :: CaoMonad m => Maybe (LExpr Name) -> m (Maybe (LExpr Var), Integer)
tcBy Nothing = return (Nothing, 1)
tcBy (Just eb) = do
    (eby', rby) <- tcIExpr RInt eb
    case rby of
        IInt n 
            | n < 0 -> return (Just (L (getLoc eby') (Lit (ILit n))), n)
            | n > 0 -> return (Just eby', n)
        _ -> tcError' SeqRangeErr

tcProc :: CaoMonad m => [LVal Name]
       -> SrcLoc
       -> SrcLoc
       -> Var
       -> [TLExpr Name]
       -> m (Stmt Var, Type Var, Class Var)
tcProc lvs fcl loc fid es = do
    (lvs', lvts, cc') <- tcLValues lvs
    let FuncSig pts frt (Proc wvars) = varType fid
    (es', pts', cc, s) <- checkArgs es pts
    let wvars' = Set.toList $ wVars (cc ++ cc') `Set.union` (Set.fromList wvars)
        frt' = applySubst s frt
        fid' = setType (FuncSig pts' frt' (Proc wvars)) fid
        expr = L fcl $ annTyE frt' $ FunCall (L loc fid') es'
    sbs <- check (zipWith (.=?>.) (fromTuple frt') lvts) []
    let lvts' = case lvts of
                    [t] -> t
                    _ -> Tuple lvts
        expr' = annotE sbs lvts' expr
    return (Assign lvs' [expr'], Bullet, Proc wvars')

tcAssign :: CaoMonad m => [LVal Name] -> [TLExpr Name] -> m (Stmt Var, Type Var, Class Var)
tcAssign lvs es = do
    (lvs', tslvs, ca) <- tcLValues lvs
    (es'', ccs) <- checkAssign es tslvs
    return (Assign lvs' es'', Bullet, maximumClass (ccs ++ ca))

checkAssign :: CaoMonad m => [TLExpr Name] -> [Type Var] -> m ([TLExpr Var], [Class Var])
checkAssign ets lvts
    | length lvts > 1 && length ets == 1 = checkTupleAssign (head ets) lvts
    | otherwise = checkMultipleAssign ets lvts

checkTupleAssign :: CaoMonad m => TLExpr Name -> [Type Var] -> m ([TLExpr Var], [Class Var])
checkTupleAssign es lvts = do
    (es', tes, c, cnstr, tp) <- tcTLExpr es
    let ets = fromTuple tes
    unless (length lvts == length ets) $ tcError' $ CardinalityException $ AssignCardinalityException TupleAssign
    sbs <- withSrcLoc (getLoc es) $ check (cnstr ++ (zipWith (.=?>.) ets lvts)) tp
    let es'' = annotE sbs (toTuple lvts) es'
    return ([es''], c)

checkMultipleAssign :: CaoMonad m => [TLExpr Name] -> [Type Var] -> m ([TLExpr Var], [Class Var])
checkMultipleAssign ets lvts
    | length lvts == length ets = zipWithAndUnzipM tcExpTyp ets lvts
    | otherwise                 = tcError' $ CardinalityException $ AssignCardinalityException MultipleAssign

-- LValue ----------------------------------------------------------------------

tcLValues :: CaoMonad m => [LVal Name] -> m ([LVal Var], [Type Var], [Class Var])
tcLValues = mapAndUnzip3M tcLValue

tcLValue :: CaoMonad m => LVal Name -> m (LVal Var, Type Var, Class Var)
tcLValue (LVVar (L loc x)) = do
    setSrcLoc loc
    (v, c) <- checkScopeLVar x
    return (LVVar (L loc v), varType v, c)
tcLValue (LVStruct lv fi) = do
    (lv', lvt, c) <- tcLValue lv
    fi' <- checkScopeSField lvt fi
    return (LVStruct lv' fi', sfType $ varType fi', c)
tcLValue (LVCont _ lv pat) = do
    (lv', lvt, c1) <- tcLValue lv
    (pat', tct, c2, cnstr) <- tcAPat lvt pat
    unless (null cnstr) $ tcError' $ UnknownErr "Non-empty constraints in a lvalue"
    return (LVCont tct lv' pat', tct, maximumClass [c1, c2])

-- XXX: Check literals despite the mode
tcAPat :: CaoMonad m => Type Var -> APat Name -> m (APat Var, Type Var, Class Var, [Constraint])
tcAPat lvt (VectP (CElem e)) = do
    (e', mi, c) <- tcAccess e
    (tct, cnstr) <- checkVBAccess lvt mi
    return (VectP (CElem e'), tct, c, cnstr)
tcAPat lvt (VectP (CRange ei ej)) = do
    (ei', i) <- tcIExpr RInt (unTypL ei)
    (ej', j) <- tcIExpr RInt (unTypL ej)
    (tct, cnstr) <- checkVBRange lvt i j
    return (VectP (CRange (annL RInt ei') (annL RInt ej')), tct, Pure, cnstr)
tcAPat lvt (MatP (CElem ei) (CElem ej)) = do
    (ei', mi, c1) <- tcAccess ei
    (ej', mj, c2) <- tcAccess ej
    (tct, cnstr) <- checkMAccess lvt (joinJust mi mj)
    return (MatP (CElem ei') (CElem ej'), tct, maximumClass [c1,c2], cnstr)
tcAPat lvt (MatP (CRange ei ej) (CRange ek el)) = do
    (ei', i) <- tcIExpr RInt (unTypL ei)
    (ej', j) <- tcIExpr RInt (unTypL ej)
    (ek', k) <- tcIExpr RInt (unTypL ek)
    (el', l) <- tcIExpr RInt (unTypL el)
    (tct, cnstr) <- checkMRange lvt i j k l
    return (MatP (CRange (annL RInt ei') (annL RInt ej')) (CRange (annL RInt ek') (annL RInt el')), tct, Pure, cnstr)
tcAPat lvt (MatP (CElem ei) (CRange ek el)) = do
    (ei', mi, c) <- tcAccess ei
    (ek', k) <- tcIExpr RInt (unTypL ek)
    (el', l) <- tcIExpr RInt (unTypL el)
    (tct, cnstr) <- checkMRow lvt k l mi
    return (MatP (CElem ei') (CRange (annL RInt ek') (annL RInt el')), tct, c, cnstr)
tcAPat lvt (MatP (CRange ek el) (CElem ei)) = do
    (ei', mi, c) <- tcAccess ei
    (ek', k) <- tcIExpr RInt (unTypL ek)
    (el', l) <- tcIExpr RInt (unTypL el)
    (tct, cnstr) <- checkMCol lvt k l mi
    return (MatP (CRange (annL RInt ek') (annL RInt el')) (CElem ei'), tct, c, cnstr)

tcAccess :: CaoMonad m => TLExpr Name -> m (TLExpr Var, Maybe (IExpr Var), Class Var)
tcAccess e = withStrictMode 
    (do (e', i) <- tcIExpr RInt (unTypL e)
        return (annL RInt e', Just i, Pure))
    (do (e', c) <- tcExpTyp e RInt
        return (e', Nothing, c))

-- VarDeclaration --------------------------------------------------------------

tcVarDecl :: CaoMonad m => Scope -> VarDecl Name -> m (VarDecl Var, Class Var)
tcVarDecl s (VarD (L loc x) b me) = do
    (b', tct) <- tcTypeDecl b
    (me', c) <- checkMExp tct me
    setSrcLoc loc
    x' <- checkDecl s tct x
    return (VarD (L loc x') b' me', c)
    where 
    checkMExp :: CaoMonad m => Type Var
              -> Maybe (TLExpr Name)
              -> m (Maybe (TLExpr Var), Class Var)
    checkMExp _ Nothing  =
        return (Nothing, Pure)
    checkMExp tct (Just e) = do
        (e', cc) <- tcExpTyp e tct
        return (Just e', cc)
tcVarDecl s (MultiD xs b) = do
    (b', tct) <- tcTypeDecl b
    xs' <- mapM (mapWithLoc (checkDecl s tct)) xs
    return (MultiD xs' b', Pure)
tcVarDecl s (ContD (L loc x) b es) = do
    (b', tct) <- tcTypeDecl b
    it <- checkContainerInit tct (genericLength es)
    (es'', ccs) <- zipWithAndUnzipM tcExpTyp es (repeat it)
    setSrcLoc loc
    x' <- checkDecl s tct x
    return (ContD (L loc x') b' es'', maximumClass ccs)

-- ConstDeclaration ------------------------------------------------------------

tcConstDecl :: CaoMonad m => Scope -> ConstDecl Name -> m (ConstDecl Var)
tcConstDecl s (ConstD (L loc x) b mo) = do
    unless globalOrLocalDef $ tcError' $ UnknownErr "Symbolic constants without value are not allowed in local definitions."
    (b', tct) <- tcIndexDecl b
    setSrcLoc loc
    (x', mo') <- checkMExp tct
    return (ConstD (L loc x') b' mo')
    where
    globalOrLocalDef = case (s, mo) of
        (Global, _) -> True
        (Local, ConstInit _) -> True
        _ -> False

    checkMExp :: CaoMonad m => Type Var -> m (Var, ConstAnn Var)
    checkMExp tct = case mo of
        None -> do
            x' <- checkConstDecl s x tct Nothing
            return (x', None)
        ConstInit me -> do
            (e, i) <- tcIExpr tct me
            x' <- checkConstDecl s x tct (Just i)
            return (x', ConstInit e)
        ConstCond me -> do
            x' <- checkConstDecl s x tct Nothing
            e <- tcCond me
            return (x', ConstCond e)

tcConstDecl Global (MultiConstD xs b me) = do
    (b', tct) <- tcIndexDecl b
    xs' <- mapM (checkVs tct) xs
    me' <- mapMaybeM tcCond me
    return (MultiConstD xs' b' me')
    where 
    checkVs :: CaoMonad m => Type Var -> Located Name -> m (Located Var)
    checkVs tct (L loc x) = 
        setSrcLoc loc >> L loc <$> checkConstDecl Global x tct Nothing

tcConstDecl Local (MultiConstD _ _ _) =
    tcError' $ UnknownErr "Symbolic constants without value are not allowed in local definitions."

tcCond :: CaoMonad m => LExpr Name -> m (LExpr Var)
tcCond e = do
    (e', i) <- tcICond e
    addHypothesis [i]
    return e'

-- Exp -------------------------------------------------------------------------

tcExpTyp :: CaoMonad m => TLExpr Name -> Type Var -> m (TLExpr Var, Class Var)
tcExpTyp e t = do
    (e', et, c, cnstr, tp) <- tcTLExpr e
    sbs <- withSrcLoc (getLoc e) $ check (cnstr ++ [et .=?>. t]) tp
    let e'' = annotE sbs t e'
    return (e'', maximumClass c)

tcTLExpr :: CaoMonad m => TLExpr Name -> m (TLExpr Var, Type Var, [Class Var], [Constraint], [TypePred])
tcTLExpr (L loc (TyE _ e)) = do
    setSrcLoc loc
    (e', tct, c, cnst, tp) <- tcExpr e
    return (L loc (annTyE tct e'), tct, c, cnst, tp)

-- Empty list means pure functions
tcExpr :: CaoMonad m => Expr Name -> m (Expr Var, Type Var, [Class Var], [Constraint], [TypePred])

tcExpr (Lit lit) = do
    (lit', tct) <- tcLiteral lit
    return (Lit lit', tct, [], [], [])

tcExpr (Var x) = do
    (x', c) <- checkScopeVar x
    return (Var x', varType x', [c], [], [])

tcExpr (FunCall (L loc fid) exps) = do
    setSrcLoc loc
    fid' <- checkScopeFunc fid
    unless (nsFunName fid') $ tcError (ScopeException fid FuncScope)
    let ft@(FuncSig pts frt c) = varType fid'
    when (isProc ft) $ tcError (BadUseException fid' ProcScope)
    (exps', pts', cs, s) <- checkArgs exps pts
    let frt' = applySubst s frt
        fid'' = setType (FuncSig pts' frt' c) fid'
    return ( FunCall (L loc fid'') exps'
            , frt'
            , c:cs
            , [], []
             )

tcExpr (StructProj sexp fi) = do
    (sexp', sexpt, c, cnstr, tp) <- tcTLExpr sexp
    fi' <- checkScopeSField sexpt fi
    return ( StructProj sexp' fi'
            , sfType $ varType fi'
            , c
            , cnstr, tp
            )

tcExpr (UnaryOp op e) = tcUnaryExpr op e

tcExpr (BinaryOp op e1 e2) = tcBinaryExpr op e1 e2

tcExpr (Access e1 pat) = do
    (e1', et1, c1, cnstr1, tp1) <- tcTLExpr e1
    (pat', tct, c2, cnstrp) <- tcAPat et1 pat
    return ( Access e1' pat'
            , tct
            , c2 : c1
            , cnstr1 ++ cnstrp
            , tp1
            )

tcExpr (Cast b td e1) = do
    (e1', et1, c, cnstr, tp) <- tcTLExpr e1
    (td', tct) <- tcTypeDeclsL td
    return (Cast b td' e1', tct, c, cnstr ++ [et1 .=>>. tct], tp)

tcUnaryExpr :: CaoMonad m => UOp -> TLExpr Name 
        -> m (Expr Var, Type Var, [Class Var], [Constraint], [TypePred])
tcUnaryExpr op e = case op of
    Sym -> do
        tct <- TyVar <$> nextTyVarId
        (e', et, c, cnstr, tp) <- tcTLExpr e
        return (UnaryOp Sym e', tct, c, cnstr ++ [ToAlgebraic tct et], Algebraic tct : tp)
    Not -> do
        (e', et, c, cnstr, tp) <- tcTLExpr e
        return (UnaryOp Not e', Bool, c,  cnstr ++ [et .=?>. Bool], tp)
    BNot -> do
        (e', et, c, cnstr, tp) <- tcTLExpr e
        return (UnaryOp BNot e', et, c, cnstr, BitsT et : tp)

tcBinaryExpr :: CaoMonad m => BinOp Name -> TLExpr Name -> TLExpr Name
        -> m (Expr Var, Type Var, [Class Var], [Constraint], [TypePred])
tcBinaryExpr bop e1 e2 = case bop of
    ArithOp op -> do
        (e1', et1, c1, cnstr1, tp1) <- tcTLExpr e1
        (e2', et2, c2, cnstr2, tp2) <- tcTLExpr e2
        (tct, cnstr', tp') <- checkAOp op et1 et2
        return ( BinaryOp (ArithOp op) e1' e2'
               , tct 
               , c1 ++ c2
               , cnstr1 ++ cnstr2 ++ cnstr'
               , tp1 ++ tp2 ++ tp'
               )
    CmpOp _ op -> do
        (e1', et1, c1, cnstr1, tp1) <- tcTLExpr e1
        (e2', et2, c2, cnstr2, tp2) <- tcTLExpr e2
        tct <- nextTyVarId
        let (cnstr, tct') = if isEqNeq op 
                then ([UnifiesC (TyVar tct) et1 et2], TyVar tct)
                else ([et1 .=?>. IntVar tct, et2 .=?>. IntVar tct], IntVar tct)
        return ( BinaryOp (CmpOp tct' op) e1' e2'
               , Bool, c1 ++ c2
               , cnstr1 ++ cnstr2 ++ cnstr
               , tp1 ++ tp2
               )
    BoolOp op -> do
        (e1', et1, c1, cnstr1, tp1) <- tcTLExpr e1
        (e2', et2, c2, cnstr2, tp2) <- tcTLExpr e2
        return (BinaryOp (BoolOp op) e1' e2'
               , Bool
               , c1 ++ c2
               , cnstr1 ++ cnstr2 ++ [et1 .=?>. Bool, et2 .=?>. Bool]
               , tp1 ++ tp2)

    BitOp op -> do
        (e1', et1, c1, cnstr1, tp1) <- tcTLExpr e1
        (e2', et2, c2, cnstr2, tp2) <- tcTLExpr e2
        return ( BinaryOp (BitOp op) e1' e2'
             , et1 -- Arbitrary choice (et1 == et2)
             , c1 ++ c2
             , cnstr1 ++ cnstr2 ++ [et1 .=?=. et2]
             , tp1 ++ tp2
             )

    BitsSROp op -> do
        (e1', et1, c1, cnstr1, tp1) <- tcTLExpr e1
        (e2', et2, c2, cnstr2, tp2) <- tcTLExpr e2
        return ( BinaryOp (BitsSROp op) e1' e2'
             , et1
             , c1 ++ c2
             , cnstr1 ++ cnstr2 ++ [et2 .=?>. RInt]
             , tp1 ++ tp2 ++ [BitsOrVector et1]
             )

    Concat -> do 
        (e1', et1, c1, cnstr1, tp1) <- tcTLExpr e1
        (e2', et2, c2, cnstr2, tp2) <- tcTLExpr e2
        tct <- TyVar <$> nextTyVarId
        return (BinaryOp Concat e1' e2'
            , tct
            , c1 ++ c2
            , cnstr1 ++ cnstr2 ++ [Conc tct et1 et2]
            , tp1 ++ tp2)

-- Literal ---------------------------------------------------------------------

-- XXX: Introduce signed literals
tcLiteral :: CaoMonad m => Literal Name -> m (Literal Var, Type Var)
tcLiteral (BLit b) =
    return (BLit b, Bool)
tcLiteral (ILit i) = do
    tvi <- IntVar <$> nextTyVarId
    return (ILit i, tvi)
tcLiteral (BSLit sig val) =
    return (BSLit sig val, Bits sig $ IInt $ genericLength val)
tcLiteral (PLit pol) = do
    (pol', tct) <- tcPolynomial pol
    return (PLit pol', tct)

-- Polynomial ------------------------------------------------------------------

tcPolynomial :: CaoMonad m => Pol Name -> m (Pol Var, Type Var)
tcPolynomial (Pol ms) = do
    (ms', mst, cnts) <- tcMonomials ms
    tct <- TyVar <$> nextTyVarId
    sbs <- check (UnifiesL tct mst : cnts) []
    let tct' = subst' sbs tct
    return (Pol ms', tct')

-- Monomial --------------------------------------------------------------------

tcMonomials :: CaoMonad m => [Mon Name] -> m ([Mon Var], [Type Var], [Constraint])
tcMonomials = fold3M tcMonomial (:) (:) (++) ([], [], [])

-- XXX: It is possible to check if the mod value is a valid one.
-- Something like 0 <= i' < mod(tct)
tcMonomial :: CaoMonad m => Mon Name -> m (Mon Var, Type Var, [Constraint])
tcMonomial (Mon (CoefI i) EZero) = do
    tct <- ModVar <$> nextTyVarId
    i' <- wellFormedExpr i
    return (Mon (CoefI i') EZero, tct, [])
tcMonomial (Mon (CoefI i) (MExpI n e)) = do
    (n', tct) <- checkScopeInd n
    i' <- wellFormedExpr i
    return (Mon (CoefI i') (MExpI n' e), tct, [])
tcMonomial (Mon (CoefP p) EZero) = do
    (p', tct) <- tcPolynomial p
    return (Mon (CoefP p') EZero, tct, [])
tcMonomial (Mon (CoefP p) (MExpI n e))= do
    (n', tct) <- checkScopeInd n
    (p', pt)  <- tcPolynomial p
    return(Mon (CoefP p') (MExpI n' e), tct, [pt .=?>. extractBaseType tct])

--------------------------------------------------------------------------------
-- Checking functions
--------------------------------------------------------------------------------

checkPolyLit :: CaoMonad m => Pol Name -> Type Var -> m ()
checkPolyLit pol tp = do
    (_, tct) <- tcPolynomial pol
    _ <- check [tct .=?>. tp] []
    return ()

checkArgs :: CaoMonad m => [TLExpr Name] -> [Type Var] -> m ([TLExpr Var], [Type Var], [Class Var], [(Var, IExpr Var)]) 
checkArgs [] [] = return ([], [], [], [])
checkArgs (e:exps) (ti:tis) = do
    (e', te, c, s) <- checkArg e ti
    (exps', tes, cs, ss) <- checkArgs exps (substAux s tis)
    return (e':exps', te:tes, c:cs, maybe ss (:ss) s)

    where
    -- XXX: inline this?
    substAux s t = maybe t (\s' -> map (Language.CAO.Syntax.Utils.subst s') t) s
-- length exps == length pts
checkArgs _ _ = tcError' $ CardinalityException ParamsCardinalityException
 
--TODO: Apply the substitution to the returned type
-- The callee function 'checkArgs' already applies the substitution to the expected type
checkArg :: CaoMonad m => TLExpr Name -> Type Var -> m (TLExpr Var, Type Var, Class Var, Maybe (Var, IExpr Var)) 
checkArg e (Index v cond ti) = do -- TODO: How to improve this?
    (e', c) <- tcExpTyp e ti
    (_, i) <- tcIExpr ti (unTypL e)
    let s = (v, i)
    case cond of
        Just cond' ->
            valid [Language.CAO.Syntax.Utils.subst s cond'] (error "<<TODO>><checkArgs>: Not satisfied condition")
        _ -> return ()
    return (e', ti, c, Just s)
    
checkArg e ti = do
    (e', c) <- tcExpTyp e ti
    return (e', ti, c, Nothing)

--------------------------------------------------------------------------------
-- TODO: Refactor to somewhere
applySubst :: [(Var, IExpr Var)] -> Type Var -> Type Var
applySubst [] t = t
applySubst (s:xs) t = applySubst xs (Language.CAO.Syntax.Utils.subst s t)