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

{- |
Module      :  $Header$
Description :  AST Post processor.
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

Post processes the CAO AST after type checking, introducing explicit casts
when coercions are used.
-}

module Language.CAO.Typechecker.PostProcessor ( annotE ) where

import Language.CAO.Typechecker.Constraint

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

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

import Language.CAO.Common.Var
import Language.CAO.Common.SrcLoc

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

annotE :: [Substitution] -> Type Var -> TLExpr Var -> TLExpr Var
annotE sbs tr (L loc (TyE ta l@(Lit _))) = let
        ta' = subst' sbs ta
    in checkCoerce (L loc $ TyE ta' l) ta' tr
annotE sbs tr (L loc (TyE ta e@(Var _))) = let
        ta' = subst' sbs ta
    in checkCoerce (L loc $ TyE ta' e) ta' tr
annotE sbs tr (L loc (TyE ta (FunCall fn es))) = let
        ta' = subst' sbs ta
        FuncSig pts _ _ = typeOf fn
        es' = zipWith (annotE sbs) pts es
    in checkCoerce (L loc $ TyE ta' $ FunCall fn es') ta' tr
-- A projection is always applied to a structure.
-- There are no coercions, unifications or casts inside strutures.
-- Structures are not algebraic types and cannot be used inside matrices
-- The expression 'se' can only be a struct or a path inside vectors or other
-- structures.
-- Conjecture: there are no undefined annotations inside the path to a structure
annotE sbs tr (L loc (TyE ta e@(StructProj _ _))) = let
        ta' = subst' sbs ta
    in checkCoerce (L loc $ TyE ta' e) ta' tr
annotE sbs tr (L loc (TyE ta (UnaryOp op e1))) = let
        ta' = subst' sbs ta
        e1' = annotE sbs ta' e1
    in checkCoerce (L loc $ TyE ta' $ UnaryOp op e1') ta' tr
annotE sbs tr (L loc (TyE ta (BinaryOp op e1 e2))) = let
        ta' = subst' sbs ta
        (e1', e2') = annotOp op sbs ta' e1 e2
    in checkCoerce (L loc $ TyE ta' $ BinaryOp op e1' e2') ta' tr
annotE sbs tr (L loc (TyE ta (Access e p))) = let
        ta' = subst' sbs ta
        e' = annotAccess sbs ta' e (isRange p)
    in checkCoerce (L loc $ TyE ta' $ Access e' p) ta' tr
-- XXX: type is being ignored
annotE sbs _ (L loc (TyE ta (Cast b td e))) = let
        ta' = subst' sbs ta
        te = subst' sbs (typeOf e)
        e' = annotE sbs te e 
    in case checkCoerce e' te ta' of
        -- Additional implicit casts may be needed
        -- Maintains the original cast external occurrence
        L _ (TyE _ (Cast _ _ e'')) -> L loc $ TyE ta' $ Cast b td e''
        -- The cast was not necessary (casting the type to itself)
        e'' -> e''
        
--------------------------------------------------------------------------------
annotOp :: BinOp Var -> [Substitution] -> Type Var -> TLExpr Var -> TLExpr Var -> (TLExpr Var, TLExpr Var)
annotOp (ArithOp Power) sbs tr e1 e2 = let
        e1' = annotE sbs tr e1
        e2' = annotE sbs Int e2
    in (e1', e2')
annotOp (ArithOp Times) sbs tr e1 e2 = let
        t1 = subst' sbs $ typeOf e1
        t2 = subst' sbs $ typeOf e2
    in case (t1, t2) of
        (Matrix s1 s2 _, Matrix s2' s3 _) -> let
                Matrix _ _ tr' = tr
                e1' = annotE sbs (Matrix s1  s2 tr') e1       
                e2' = annotE sbs (Matrix s2' s3 tr') e2
            in (e1', e2')
        _ -> let
                e1' = annotE sbs tr e1       
                e2' = annotE sbs tr e2
            in (e1', e2')
annotOp (CmpOp ty _) sbs Bool e1 e2 = let
        tr = subst' sbs ty
        e1' = annotE sbs tr e1
        e2' = annotE sbs tr e2
    in (e1', e2')
annotOp (CmpOp _ _) _ _ _ _ = error "<PostProcessor>.<annotOp>: Unexpected case for compare operations"
annotOp (BitsSROp _) sbs tr e1 e2 = let
        e1' = annotE sbs tr e1
        e2' = annotE sbs RInt e2
    in (e1', e2')
annotOp Concat sbs tr e1 e2 = case tr of
    -- Concatenation of bit strings does not imply any coercion
    Bits _ _ -> let
            t1 = subst' sbs $ typeOf e1
            t2 = subst' sbs $ typeOf e2
            e1' = annotE sbs t1 e1
            e2' = annotE sbs t2 e2
        in (e1', e2')
    Vector _ t -> let
            -- Maintain the sizes of the expected vector types
            Vector k1 _ = subst' sbs $ typeOf e1
            Vector k2 _ = subst' sbs $ typeOf e2
            e1' = annotE sbs (Vector k1 t) e1
            e2' = annotE sbs (Vector k2 t) e2
        in (e1', e2')
    _ -> error "<PostProcessor>.<annotOp>: Unexpected case for concat"
annotOp _ sbs tr e1 e2 = let
        e1' = annotE sbs tr e1       
        e2' = annotE sbs tr e2
    in (e1', e2')

--------------------------------------------------------------------------------
-- In this function, we have to distinguish the cases of accesses to individual 
-- elements, from the case of ranges.
annotAccess :: [Substitution] -> Type Var -> TLExpr Var -> Bool -> TLExpr Var
annotAccess sbs tr e isRng = case subst' sbs (typeOf e) of
    t@(Bits _ _) -> annotE sbs t e
    Vector k _ -> annotE sbs (Vector k tr') e
    Matrix i j _ -> annotE sbs (Matrix i j tr') e
    _ -> error "<PostProcessor>.<annotOp>: Unexpected case for an access"
    where 
    tr' = if isRng then head (innerType tr) else tr

--------------------------------------------------------------------------------
checkCoerce :: TLExpr Var -> Type Var -> Type Var -> TLExpr Var
checkCoerce e tct tct'
    | tct == tct'                   = e
    | isModInt tct && isModInt tct' = 
        let e' = addCoerce e Int
        in checkCoerce e' Int tct'
    | isModInt tct && isBits tct'   = 
        let e' = addCoerce e Int
        in checkCoerce e' Int tct'
    | isBits tct && isMod tct'      =
        let e' = addCoerce e Int
        in checkCoerce e' Int tct'
    | isInt tct && isModPol tct'    =
        let bas = extractBottomBaseType tct'
            e' = addCoerce e bas
        in checkCoerce e' bas tct'
    | isMod tct && isModPol tct' && extractBaseType tct' /= tct =
        let bas = extractBaseType tct'
            e' = addCoerce e bas
        in checkCoerce e' bas tct'
    | otherwise
        = addCoerce e tct'

--------------------------------------------------------------------------------
addCoerce :: TLExpr Var -> Type Var -> TLExpr Var
addCoerce e tct' = let 
        tct = case tct' of
            Tuple tpl -> map (L genSrcLoc . type2TyDecl) tpl
            _ -> [L genSrcLoc $ type2TyDecl tct']
    in L genSrcLoc $ annTyE tct' $ Cast False tct e

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