{- |
Module      :  $Header$
Description :  Type checking indexes
Copyright   :  (c) SMART Team / HASLab
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

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