{-  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/>. -}

{- |
Module      :  $Header$
Description :  Type checking indexes
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
-}

module Language.CAO.Typechecker.Index 
    ( tcIndexDecl
    , tcICond
    , tcIExpr
    , wellFormedExpr
    ) where

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

import Language.CAO.Common.Error
import Language.CAO.Common.Literal
import Language.CAO.Common.Monad
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.Index.Eval
import Language.CAO.Index.Utils
import Language.CAO.Syntax
import Language.CAO.Typechecker.Check
import Language.CAO.Typechecker.Expr
import Language.CAO.Typechecker.PostProcessor
import Language.CAO.Type
import Language.CAO.Type.Utils

--------------------------------------------------------------------------------
-- Type checking index declarations

tcIndexDecl :: CaoMonad m => TyDecl Name -> m (TyDecl Var, Type Var)
tcIndexDecl IntD  = return (IntD, Int)
tcIndexDecl RIntD = return (RIntD, RInt)
tcIndexDecl BoolD = return (BoolD, Bool)
tcIndexDecl _ = tcError (NotSupportedIndexTyp :: ErrorCode Var)
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
-- Conditions

tcICond :: CaoMonad m => LExpr Name -> m (LExpr Var, ICond Var)
tcICond (L loc e) = 
    setSrcLoc loc >> mapPair (L loc) evalCond <$> tceICond' e

tceICondTL :: CaoMonad m => TLExpr Name -> m (TLExpr Var, ICond Var)
tceICondTL (L loc (TyE _ e)) = withSrcLoc loc $
    mapFst (L loc . annTyE Bool) <$> tceICond' e

tceICond' :: CaoMonad m => Expr Name -> m (Expr Var, ICond Var)
tceICond' (Lit (BLit b)) = 
    return (Lit (BLit b), IBool b)
tceICond' (Var x) = checkIndex x
    (ifM (isBool . varType) (return . split Var IBInd) 
                            (const $ tcError (NotSupportedIndexOp :: ErrorCode Var)))
    (const $ const $ tcError (NotSupportedIndexOp :: ErrorCode Var))
tceICond' (UnaryOp Not e) = 
    mapPair (UnaryOp Not) INot <$> tceICondTL e

tceICond' (BinaryOp (CmpOp _ op) e1 e2) = do
    (e1', i1, et1, c1) <- tieTL e1
    (e2', i2, et2, c2) <- tieTL e2
    tct <- TyVar <$> nextTyVarId
    sbs <- check (c1 ++ c2 ++ [Unifies tct et1 et2]) []
    let tct' = subst' sbs tct
        e1'' = annotE sbs tct' e1'
        e2'' = annotE sbs tct' e2'
    return (BinaryOp (CmpOp tct' op) e1'' e2'', (mapCOp op) i1 i2)

tceICond' (BinaryOp (BoolOp op) e1 e2) = do
    (e1', i1) <- tceICondTL e1
    (e2', i2) <- tceICondTL e2
    let i = case op of
            And -> IAnd [i1, i2]
            _   -> (mapBOp op) i1 i2
    return (BinaryOp (BoolOp op) e1' e2', i)

tceICond' _ = tcError (NotSupportedIndexOp :: ErrorCode Var)

--------------------------------------------------------------------------------
-- Expressions

tcIExpr :: CaoMonad m => Type Var -> LExpr Name -> m (LExpr Var, IExpr Var)
tcIExpr tpr (L loc e) =
    setSrcLoc loc >> mapPair (L loc) evalExpr <$> topTie tpr e

topTie :: CaoMonad m => Type Var -> Expr Name -> m (Expr Var, IExpr Var)
topTie tpr e = do
    (e', i, t, c) <- tie e
    sbs <- check (c ++ [t .=?>. tpr]) []
    let e'' = unTyp $ unLoc $ annotE sbs tpr $ genLoc $ annTyE t e'
    return (e'', i)

tieTL :: CaoMonad m => TLExpr Name -> m (TLExpr Var, IExpr Var, Type Var, [Constraint])
tieTL (L loc (TyE _ e)) = withSrcLoc loc $ do
    (e', i, t, c) <- tie e
    return (L loc (annTyE t e'), i, t, c)

tie :: CaoMonad m => Expr Name -> m (Expr Var, IExpr Var, Type Var, [Constraint])
tie (Lit (ILit i)) = do
    tid <- IntVar <$> nextTyVarId
    return (Lit (ILit i), IInt i, tid, [])
tie (Var x) = do
    (x', i) <- checkIndex x
        (ifM (isInt . varType)
             (\ x' -> return $ maybe (x', IInd x') (curry id x') $ indConst x')
             (const $ tcError (NotSupportedIndexOp :: ErrorCode Name)))
        (return .$. curry (mapSnd IInt))
    return (Var x', i, varType x', [])
tie (BinaryOp (ArithOp op) e1 e2) = do
    (e1', i1, et1, c1) <- tieTL e1
    (e2', i2, et2, c2) <- tieTL e2
    tct <- TyVar <$> nextTyVarId
    let iexp = case op of
            Plus -> ISum [i1, i2]
            _ -> (mapAOp op) i1 i2
        cnst = [Unifies tct et1 et2]
    return (BinaryOp (ArithOp op) e1' e2', iexp, tct, c1 ++ c2 ++ cnst)
tie (UnaryOp Sym e) = do
    (e', i, et, c) <- tieTL e
    return (UnaryOp Sym e', ISym i, et, c)

tie _ = tcError (NotSupportedIndexOp :: ErrorCode Name)
    
--------------------------------------------------------------------------------
wellFormedExpr :: CaoMonad m => IExpr Name -> m (IExpr Var)
wellFormedExpr = liftM evalExpr . worker
    where
    worker ie = case ie of
        IInt n          -> return $ IInt n
        IArith op i1 i2 -> liftM2 (IArith op) (worker i1) (worker i2)
        ISym i          -> ISym <$> worker i
        ISum l          -> ISum <$> mapM worker l
        IInd _          -> return (undefined) -- TODO: verify if it is in environment

--------------------------------------------------------------------------------