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

{- |
Module      :  $Header$
Description :  
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.Check (
      checkScopeDepIndex
    , checkScopeInd
    , checkScopeConst
    , checkScopeConst'
    , checkIndex
    , checkScopeType
    , checkScopeFunc
    , checkScopeProc
    , checkScopeVar
    , checkScopeLVar
    , checkScopeSField

    , checkFuncReturn
    , checkPolynomial
    , checkContainerInit
    , checkAOp

    , checkDecl
    , checkConstDecl

    , checkTySyn
    , checkStructDecl
    , checkBitsSize
    , checkMod
    , checkModBase
    
    , checkVectorSize
    , checkMatrixSize
    , checkAlgebraic

    , checkVBAccess
    , checkVBRange
    , checkMAccess
    , checkMRange
    , checkMRow
    , checkMCol

    ) where

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

import Data.Maybe ( fromJust )

import Language.CAO.Common.Outputable (PP)

import Language.CAO.Common.Error
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.Index.Eval

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

import Language.CAO.Typechecker.Constraint
import Language.CAO.Typechecker.Expr
import Language.CAO.Typechecker.Heap
import Language.CAO.Typechecker.Solver

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

--------------------------------------------------------------------------------
-- TODO: this functions are instances of the same schema
-- TODO: Some use containsXXX other lookupXXX but they are
-- the same thing...
genericCheck :: CaoMonad m => (Name -> Int -> Type Var -> Var) 
             -> (Heap -> Var -> Heap)
             -> Type Var
             -> Name
             -> m Var
genericCheck f fh t x = getHeap >>= \ h -> do
    checkContains h x
    u <- newUniq
    let v = f x u t
    putHeap $ fh h v
    return v

checkContains :: CaoMonad m => Heap -> Name -> m ()
checkContains h x =
    when (containsName h x) $ tcError $ DeclException $ MultipleDeclException $ fromJust $ lookupName h x

genericScope :: CaoMonad m => Name -> (Var -> m a) -> (Var -> m a) -> m a
genericScope x fGlobal fLocal = getHeap >>= \h ->
    case lookupLocalName h x of
        Nothing -> case lookupGlobalName h x of
            Nothing -> tcError $ ScopeException x VarScope
            Just v -> fGlobal v
        Just v -> fLocal v

checkGlobal :: (PP id, Show id, Read id, CaoMonad m) => Name -> (Var -> m a) -> ErrorCode id -> m a
checkGlobal x fjust err = getHeap >>= 
    maybe (tcError err) fjust . flip lookupGlobalName x

varCheck :: CaoMonad m => (Var -> Bool) -> (Var -> a) -> Var -> m a
varCheck cond modifier = ifM cond (return . modifier) checkBadUse

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

checkDecl :: CaoMonad m => Scope -> Type Var -> Name -> m Var
checkDecl Local  = genericCheck mkLId insertLocalName
checkDecl Global = genericCheck mkGId insertGlobalName

checkConstDecl :: CaoMonad m => Scope -> Name -> Type Var -> Maybe (IExpr Var) -> m Var
checkConstDecl scope x t e = case scope of
    Local ->  genericCheck (mkConst mkLConst) insertLocalName  t x
    Global -> genericCheck (mkConst mkGConst) insertGlobalName t x
    where
    mkConst f x' u' t' = f x' u' t' e

checkTySyn :: CaoMonad m => Name -> Type Var -> m Var
checkTySyn x t = genericCheck aux insertGlobalName t x
    where
    aux x' u' t' = let
            v = mkGId x' u' tct
            tct = TySyn v t'
        in v

--------------------------------------------------------------------------------
checkStructDecl :: CaoMonad m => Located Name
                -> [Located Name]
                -> [TyDecl Var]
                -> [Type Var]
                -> m (TyDef Var)
checkStructDecl ln@(L loc n) flds tds tys = getHeap >>= \ h -> do
    mapM_ (checkContains h . unLoc) $ ln:flds
    u  <- newUniq
    us <- mapM (const newUniq) flds
    let n'     = mkGId n u t'
        t'     = TySyn n' $ Struct n' fts
        flds'  = map (\ (L lc vv, uu, tt) ->
                         L lc $ mkGId vv uu (SField n' tt)) vvuutt
        fts    = zip (map unLoc flds') tys
        ftds   = zip flds' tds
        vvuutt = zip3 flds us tys
  -- XXX: change to a strict version of foldl
    putHeap $ foldl insertGlobalName h (n': map unLoc flds')
    return (StructDecl (L loc n') ftds)

checkPolynomial :: CaoMonad m => Type Var -> Name -> Pol Name -> m (Var, Pol Var, Type Var)
checkPolynomial td ti pol = getHeap >>= \ h -> 
    case lookupGlobalName h ti of
        Nothing -> do 
            u <- newUniq
            let v = mkGId ti u (Indet t)
                t    = Mod (Just td) (Just v) pol'
                pol' = pol <|> ti ~> v
            putHeap (insertGlobalName h v)
            return (v, pol', t)
        Just v -> do 
            let pol' = pol <|> ti ~> v
                ty   = Mod (Just td) (Just v) pol'
            unless (varType v == Indet ty) $ tcError $ DeclException $ MultipleDeclException v
            return (v, pol', ty)
        
--------------------------------------------------------------------------------
checkScopeVar :: CaoMonad m => Name -> m (Var, Class Var)
checkScopeVar xvar = genericScope xvar fGlobal fLocal
    where
    fGlobal v
        | nsVar v   = return (v, RO)
        | indVar v  = return (v, Pure)
        | otherwise = checkBadUse v
    fLocal v
        | nsVar v   = return (v, Pure)
        | indVar v  = return (v, Pure)
        | otherwise = checkBadUse v

checkScopeDepIndex :: CaoMonad m => Name -> m Var
checkScopeDepIndex ind = genericScope ind fIndex fIndex
    where
    fIndex = varCheck indVar id

checkScopeLVar :: CaoMonad m => Name -> m (Var, Class Var)
checkScopeLVar lvar = genericScope lvar fGlobal fLocal
    where
    fGlobal = varCheck nsVar (split id (Proc . singleton))
    fLocal  = varCheck nsVar (split id (const Pure))

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

checkScopeConst' :: CaoMonad m => Name -> m (Maybe Integer)
checkScopeConst' cnst = getHeap >>= return . flip lookupConstName cnst

checkScopeConst :: CaoMonad m => Name -> m Integer
checkScopeConst cnst = 
    getHeap >>= maybe (tcError (IntEvalErr :: ErrorCode Name)) return . flip lookupConstName cnst

checkIndex :: CaoMonad m => Name -> (Var -> m a) -> (Var -> Integer -> m a) -> m a
checkIndex x fnothing fjust = do
    x' <- checkScopeDepIndex x
    mi <- checkScopeConst' $ varName x
    maybe (fnothing x') (fjust x') mi
    
--------------------------------------------------------------------------------

checkScopeType :: CaoMonad m => Name -> m Var
checkScopeType syn = checkGlobal syn
    (varCheck (isTySyn . varType) id)
    (ScopeException syn TypeScope)

checkScopeProc :: CaoMonad m => Name -> m Var
checkScopeProc fid = checkGlobal fid
    (varCheck (isProc . varType) id)
    (ScopeException fid ProcScope)

checkScopeFunc :: CaoMonad m => Name -> m Var
checkScopeFunc fid = checkGlobal fid
    (varCheck nsFunName id)
    (ScopeException fid FuncScope)

checkScopeInd :: CaoMonad m => Name -> m (Var, Type Var)
checkScopeInd indx = checkGlobal indx
    (\ t -> case varType t of
                Indet ty -> return (t, ty)
                _        -> checkBadUse t)
    (ScopeException indx IndScope)

checkScopeSField :: CaoMonad m => Type Var -> Name -> m Var
checkScopeSField (Struct sn1 _) fi = checkGlobal fi
    (\ v -> case varType v of
                SField sn2 _ | sn1 == sn2 -> return v
                _  -> tcError $ ScopeException sn1 (SFieldScope fi))
    (ScopeException sn1 (SFieldScope fi))
checkScopeSField st _ = tcError (WrongTypeException st StructType)

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

checkBadUse :: CaoMonad m => Var -> m a
checkBadUse x = tcError $ BadUseException x $ checkAux $ varType x
    where
    checkAux t
        | isProc t     = ProcScope
        | isFunType t  = FuncScope
        | isVar t      = VarScope
        | isTySyn t    = TypeScope
        | isIndet t    = IndetScope
        | otherwise    = GenericScope

checkFuncReturn :: CaoMonad m => Type Var -> Type Var -> m ()
checkFuncReturn t1 rt
  = unless (not (isNil t1) || isNil rt) $ tcError (FuncReturnErr :: ErrorCode Var)

checkContainerInit :: CaoMonad m => Type Var -> Integer -> m (Type Var)
checkContainerInit (Vector k it) n = do
    valid [ k .==. IInt n ] $ CardinalityException $ InitCardinalityException VectorType
    return it
checkContainerInit (Matrix u v it) n = do
    checkAlgebraic it
    valid [ (u .*. v) .==. IInt n ] $ CardinalityException $ InitCardinalityException VectorType
    return it
checkContainerInit _ _ = tcError (ContainerInitErr :: ErrorCode Var)

checkBitsSize :: CaoMonad m => IExpr Var -> m ()
checkBitsSize s = valid [IInt 1 .<=. s] (DeclException (SizeDeclException s Nothing BitsType))
    
-- Precondition: The index arguments are reduced
checkVectorSize :: CaoMonad m => IExpr Var -> m ()
checkVectorSize s = valid [ IInt 1 .<=. s ] (DeclException (SizeDeclException s Nothing VectorType))

-- Precondition: The index arguments are reduced
checkMatrixSize :: CaoMonad m => IExpr Var -> IExpr Var -> m ()
checkMatrixSize r c = valid [ IInt 1 .<=. r, IInt 1 .<=. c ] (DeclException (SizeDeclException r (Just c) MatrixType))

checkModBase :: CaoMonad m => IExpr Var -> m ()
checkModBase b = valid [IInt 2 .<=. b] (DeclException $ BaseDeclException b)

-- By the standard, the bottom base type has to be a mod!
checkMod :: CaoMonad m => Type Var -> m ()
checkMod t = unless (isMod t && isMod (extractBottomBaseType t)) $ tcError (WrongTypeException t ModType)

--------------------------------------------------------------------------------
-- TODO: The verification of accesses is very similar to the code of unification
checkMAccess :: CaoMonad m => Type Var -> Maybe (IExpr Var, IExpr Var) -> m (Type Var, [Constraint])
checkMAccess (Matrix u v it) mi = do
    cAccessM u v mi
    return (it, [])
checkMAccess t mi = do
    tid <- TyVar <$> nextTyVarId
    return (tid, [MAccess tid t mi])

cAccessM :: CaoMonad m => IExpr Var -> IExpr Var -> Maybe (IExpr Var, IExpr Var) -> m ()
cAccessM _ _ Nothing = return ()
cAccessM u v (Just (i, j)) = 
    valid [IInt 0 .<=. i, i .<. u, IInt 0 .<=. j, j .<. v] $ UnknownErr "Checking strict access (matrix access)"

checkMRange :: CaoMonad m => Type Var -> IExpr Var -> IExpr Var -> IExpr Var -> IExpr Var -> m (Type Var, [Constraint])
checkMRange (Matrix u v it) i j k l = do
    uu <- checkRange u i j (RangeException MatrixType)
    vv <- checkRange v k l (RangeException MatrixType) 
    return (Matrix uu vv it, [])
checkMRange t i j k l = do
    tid <- TyVar <$> nextTyVarId
    return (tid, [MRange tid t i j k l])

checkMRow :: CaoMonad m => Type Var -> IExpr Var -> IExpr Var -> Maybe (IExpr Var) -> m (Type Var, [Constraint])
checkMRow (Matrix v u it) i j ma = do
    cAccessV v ma
    vv <- checkRange u i j (RangeException MatrixType)
    return (Matrix (IInt 1) vv it, [])
checkMRow t i j ma = do
    tid <- TyVar <$> nextTyVarId
    return (tid, [MRow tid t i j ma])

checkMCol :: CaoMonad m => Type Var -> IExpr Var -> IExpr Var -> Maybe (IExpr Var) -> m (Type Var, [Constraint])
checkMCol (Matrix v u it) i j ma = do
    cAccessV u ma
    vv <- checkRange v i j (RangeException MatrixType)
    return (Matrix vv (IInt 1) it, [])
checkMCol t i j ma = do
    tid <- TyVar <$> nextTyVarId
    return (tid, [MCol tid t i j ma])
    

checkVBAccess :: CaoMonad m => Type Var -> Maybe (IExpr Var) -> m (Type Var, [Constraint])
checkVBAccess (Bits s k) i = do
    cAccessV k i
    return (Bits s (IInt 1), [])
checkVBAccess (Vector k it) i = do
    cAccessV k i
    return (it, [])
checkVBAccess t1 i = do
    tid <- TyVar <$> nextTyVarId
    return (tid, [VBAccess tid t1 i])

cAccessV :: CaoMonad m => IExpr Var -> Maybe (IExpr Var) -> m ()
cAccessV _ Nothing = return ()
cAccessV k (Just i) = 
    valid [IInt 0 .<=. i, i .<. k] $ UnknownErr "Checking strict access (vector)"
checkVBRange :: CaoMonad m => Type Var -> IExpr Var -> IExpr Var -> m (Type Var, [Constraint])
checkVBRange (Bits s k) i j     = do
    k' <- checkRange k i j (RangeException BitsType)
    return (Bits s k', [])
checkVBRange (Vector k it) i j = do
    k' <- checkRange k i j (RangeException VectorType)
    return (Vector k' it, [])
checkVBRange t1 i j = do
    tid <- TyVar <$> nextTyVarId
    return (tid, [VBRange tid t1 i j])

checkRange :: CaoMonad m => IExpr Var -> IExpr Var -> IExpr Var -> ErrorCode Var -> m (IExpr Var)
checkRange u i j err = do
    valid [j .<. u, i .<=. j, IInt 0 .<=. i] err
    return $ evalExpr $ ISum [ j, ISym i, IInt 1 ]

--------------------------------------------------------------------------------
checkAOp :: CaoMonad m => AOp -> Type Var -> Type Var -> m (Type Var, [Constraint], [TypePred])
checkAOp Times t1 t2 = do
    tct <- TyVar <$> nextTyVarId
    return (tct, [Mult tct t1 t2], [Algebraic tct])
checkAOp Power t1 t2 = do
    tct <- TyVar <$> nextTyVarId
    return (tct, [Pow tct t1, t2 .=?>. Int], [Algebraic tct])
checkAOp Div   t1 t2 = do
    tct <- TyVar <$> nextTyVarId
    return (tct, [Unifies tct t1 t2], [IntOrMod tct])
checkAOp ModOp t1 t2 = do
    tct <- IntVar <$> nextTyVarId
    return (tct, [t1 .=?>. tct, t2 .=?>. tct], [])
-- TODO: This does not work for operation on bit strings
checkAOp _     t1 t2 = do
    tct <- TyVar <$> nextTyVarId
    return (tct, [Unifies tct t1 t2], [Algebraic tct])