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

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

The simplification step aims at reducing the mismatch between CAO and C or, 
more precisely, the C backend. Compilers that generate assembly code 
traditionally use an intermediate representation known as three-address code,
in which every instruction is in its simpler form with two operand addresses
and one result address. Since we are targeting a software API as a backend, our
format is quite different but shares some of the same principles. Operations 
in the backend expect variables or constants as parameters and a variable as 
result. This means that nested expressions must be extracted and replaced by 
an auxiliary variable. For instance, the following assignment of an arithmetic 
expression:

> a := 3 * b + 2 * c - 4;

should be transformed to:

@
def t0 : int;
def t1 : int;
def t2 : int;
t0 := 3  * b;
t1 := 2  * c;
t2 := t0 + t1;
a  := t2 - 4;
@

In general, for assignments of results from binary and unary operations, we 
must obey the following format in which @op2@ is a binary operator, @op1@ is an 
unary operator, @var@ is a variable and @e1@ and @e2@ are either variables or 
constants:

@
     :=
    /  \
 var   op2
      /  \
    e1    e2
@

@
     :=
    /  \
 var   op1
        |
        e1
@

Similar formats were defined for the other operations, leading to a normalized
code format, ready to be translated to C. This normalization process also 
includes the following actions:

* The initialization of global variables is removed from their declaration and
put in a global init procedure.

* Simultaneous variable declarations are transformed to multiple simple variable
declarations.

* The initializations of variable declarations are removed from declarations 
and added as independent statements, except for container initializations 
(vectors and matrices). Here, the natural way of simplifying this kind of 
initializations would be a position-wise assignment. However, this would 
invalidate block initialization during translation. Therefore, only expressions
inside container initializations are simplified.

* Parallel assignments are transformed to several simple assignments 
(except parallel assignments from multiple function results which can only be 
resolved during translation with the introduction of references).

* All other expressions are simplified in order that operands may be either 
variables or constants.

-}
module Language.CAO.Transformation.Simplify (
    simplifyCaoAST
 ) where

import Control.Monad.State
import Data.DList ()
import qualified Data.DList as DL
import qualified Data.List as Lst
import Data.Maybe
import Data.Set ()
import qualified Data.Set as Set

import Language.CAO.Common.Error
import Language.CAO.Common.Fresh
import Language.CAO.Common.Monad
import Language.CAO.Common.Polynomial
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.Utils
import Language.CAO.Syntax
import Language.CAO.Syntax.Utils
import Language.CAO.Type
import Language.CAO.Type.Utils

--
--The left value code is not correctly handled during the translation:
-- seq i := 1 to 30 {
--     vec[i+3-n][n*3*i+i] := vec[i+3-n][n*3*i+i];
-- }
--
--Without renaming this temporary variable generation schema fails if there
--is any identifier of the form "t0", "t1", ... on the code

type SS = LStmt Var
type ConstDef = LStmt Var
type VarDeclaration = LStmt Var

-- CaoAST ----------------------------------------------------------------------

-- | Applies the simplification step to the AST. Takes as parameter, the name
--   of the global inititialization procedure.
simplifyCaoAST :: CaoMonad m => String -> Prog Var -> m (Prog Var)
simplifyCaoAST initProcName (Prog defs _) = withSimplifyST $ do
    (defs', stmts, cdecl, vdecl) <- concatMapAndUnzip4M simplifyDef defs
    -- The global initialization procedure is only necessary if there is
    -- something to initialize.
    let initDef = if null stmts && null cdecl && null vdecl
            then Nothing
            else let 
             -- Written global variables inside the init procedure, i.e.,
             -- initialization of global variables.
             -- Testing for assignments is necessary [See note 8]
             wvars   = Set.toList $ Set.filter isGlobalVar $ 
                        fvs $ filter (isAssignStmt . unLoc) stmts
             fName   = globalInit initProcName wvars
             body    = funcBody cdecl vdecl stmts
             in Just $ Fun (genLoc fName) [] [] body
    return $ Prog defs' initDef

-- Definition ------------------------------------------------------------------
{-
This function returns:
* A list of global definitions. Since multiple variable definitions are 
transformed into several individual declarations, as list is needed.
* A list of assignments used as initialization of global declarations in the
global initialization procedure. We should notice that besides assign 
statements, also container declarations are used [See note 8].
* A list of auxiliary constant definitions (declaration + initialization)
to be used in the global initialization procedure.
* A list of auxiliary variable declarations to be used in the global 
initialization procedure.

-}
simplifyDef 
    :: CaoMonad m
    => LDef Var 
    -> m ( [LDef Var]
         , [SS]
         , [ConstDef]
         , [VarDeclaration]
         )
simplifyDef (L l (VarDef vd)) = do
    (vd', stmts, index, decl) <- simplifyVarDeclaration vd
    return (map (L l . VarDef) vd', stmts, index, decl)
simplifyDef (L l (FunDef f)) = do
    f' <- simplifyFunc f
    return (L l (FunDef f') : [], [], [], [])
simplifyDef d@(L _ (TyDef _)) = return (d : [], [], [], [])
simplifyDef (L l (ConstDef cd)) = do
    (cd', index) <- simplifyConstDeclaration cd True
    return (L l (ConstDef cd') : [], [], DL.toList index, [])


-- Func ------------------------------------------------------------------------
{- 
Note 1:
The introduction of depedent types implies that type expression have to be
simplified, too. For instance, the index of type:
@vector[3 * n + 1] of int@
must be broken down to simple expressions:

@
t0 := 3 * n;
t1 := t0 + 1;
vector [t1] of int;
@

To simplify this process, a environment was introduced to hold new types for
variables. This means that the simplification process is done only once and
then all uses of the same variable are immediately retyped. This environment
has to be reset everytime a function body is processed.
-}
simplifyFunc :: CaoMonad m =>
    Fun Var -> m (Fun Var)
simplifyFunc (Fun fname args rtype body) = do
    resetSimplifyST -- [See Note 1]
    (body', index, decl) <- simplifyStatements body
    let body'' = funcBody' index decl body'
    return (Fun fname args rtype body'')

{-
Note 2: 
The order of statements inside a function is important since there are
dependencies between declarations, definitions and assignments. Thus, the
body is divided in logical blocks, marked by annotations. The overall schema
is the following:

[ Index (simbolic variable) declarations ]
  Nop EndIndex
[ Auxiliary variable declarations ]
  Nop EndAux
[ Assignemnts to variables ]

This schema is needed because 
- indexes may be used in the declaration of variables.
- subsequent phases of the compiler pipeline must know where each block ends.
(better explain this point).
-}
funcBody :: [ConstDef] -> [VarDeclaration] -> [SS] -> [LStmt Var]
funcBody index decl body = fBody (++) (:) index decl body

funcBody' :: DL.DList ConstDef 
          -> DL.DList VarDeclaration 
          -> DL.DList SS -> [LStmt Var]
funcBody' index decl body = DL.toList $
    fBody DL.append DL.cons index decl body

fBody :: (t -> t1 -> t2) 
      -> (Located (Stmt id) -> t2 -> t1) 
      -> t -> t -> t2 -> t2
fBody append cons index decl body = 
    index 
    `append` 
    ((genLoc $ Nop EndIndex) `cons` (decl 
    `append` 
    ((genLoc $ Nop EndAux) `cons` body)))

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

-- For statements, this is mostly traversal code.

simplifyStatements 
    :: CaoMonad m 
    => [LStmt Var] 
    -> m (DL.DList (LStmt Var), DL.DList ConstDef, DL.DList VarDeclaration)
simplifyStatements = concatMapAndUnzip3MD simplifyStatement

simplifyStatement 
    :: CaoMonad m 
    => LStmt Var 
    -> m (DL.DList (LStmt Var), DL.DList ConstDef, DL.DList VarDeclaration)
simplifyStatement (L l s) = simplifyStmt l s

simplifyStmt 
    :: CaoMonad m 
    => SrcLoc -> Stmt Var 
    -> m (DL.DList (LStmt Var), DL.DList ConstDef, DL.DList VarDeclaration)

simplifyStmt l (VDecl vd) = simplifyLocalVarDeclaration l vd
simplifyStmt l (CDecl cd) = do
    (cd', cdecls) <- simplifyConstDeclaration cd False
    return (DL.empty, cdecls `DL.snoc` L l (CDecl cd') , DL.empty)

simplifyStmt l (Assign lv' e') =
    case (lv', e') of
        (lv:[], e:[] ) -> simplifyAssignment l lv e
        (_:_:_, e:[] ) -> simplifyTupleAssignment l lv' e
        (_:_:_, _:_:_) -> simplifyMultipleAssignment l lv' e'
        _              -> error "<Language.CAO.Transformation.Simplify>.\
            \<simplifyStmt>: unexpected case in assignment"

-- XXX: Is it necessary to update type annotations?
simplifyStmt l (FCallS fid exps) = do
    (exps', stmts, cdecl, vdecl) <- simplifyExps exps
    return (stmts `DL.snoc` L l (FCallS fid (DL.toList exps')), cdecl, vdecl)

simplifyStmt l (Ret exps) = do
    (exps', stmts, cdecl, vdecl) <- simplifyExps exps
    return (stmts `DL.snoc` L l (Ret (DL.toList exps')), cdecl, vdecl)

simplifyStmt l (Ite i t e) = do
    (cond', stmts, cdecl1, vdecl1) <- simplifyExpChoice i
    (i', cdecl2, vdecl2)           <- simplifyStatements t
    (e', cdecl3, vdecl3)           <- simplifyM e
    return ( stmts  `DL.snoc`   L l (Ite cond' (DL.toList i') e')
           , cdecl1 `DL.append` cdecl2 `DL.append` cdecl3
           , vdecl1 `DL.append` vdecl2 `DL.append` vdecl3)
    where 
    simplifyM Nothing  = return (Nothing, DL.empty, DL.empty)
    simplifyM (Just s) = do
        (e'', cdecl, vdecl) <- simplifyStatements s
        return (Just (DL.toList e''), cdecl, vdecl)

simplifyStmt l (While cond wstms) = do
    (cond', stmts, cdecl1, vdecl1) <- simplifyExpChoice cond
    (wstms', cdecl2, vdecl2) <- simplifyStatements wstms
    -- The condition has to be added to the end of the body [See Note 3]
    let wbody = DL.toList $ wstms' `DL.append` stmts
    return ( stmts  `DL.snoc`   L l (While cond' wbody)
           , cdecl1 `DL.append` cdecl2
           , vdecl1 `DL.append` vdecl2)

--- XXX: type annotations
-- The simplification of the bound values can make pointer from integers
-- depending of the backend
simplifyStmt l (Seq (SeqIter v s e b r) sstms) = do
    -- A type annotation was added because 'simplifyExpChoice' expects it.
    -- The bounds must always be of type RInt
    (s', st1, cdecl1, vdecl1) <- simplifyExpChoice (annL RInt s)
    (e', st2, cdecl2, vdecl2) <- simplifyExpChoice (annL RInt e)
    (b', st3, cdecl3, vdecl3) <- simplifyM b
    (sstms' , cdecl4, vdecl4) <- simplifyStatements sstms
    -- All variables/constants that do not depend on the index may be declared
    -- only once outside the body of the sequence.
    -- Otherwise, they have to be declared inside the body of the function
    let (cdeclS, cdeclO) = innerConsts v $ DL.toList cdecl4 -- [See Note]
        (vdeclS, vdeclO) = innerVars (v : declaredConsts cdeclS) $ DL.toList vdecl4
        sstms'' = cdeclS ++ vdeclS ++ DL.toList sstms'
        cdeclO' = DL.fromList cdeclO
        vdeclO' = DL.fromList vdeclO
    return ( st1    `DL.append` st2    `DL.append` st3    `DL.snoc` 
        L l ( Seq (SeqIter v (unTypL s') (unTypL e') b' r) sstms'' )
           , cdecl1 `DL.append` cdecl2 `DL.append` cdecl3 `DL.append` cdeclO'
           , vdecl1 `DL.append` vdecl2 `DL.append` vdecl3 `DL.append` vdeclO')
    where 
    simplifyM Nothing   = return (Nothing, DL.empty, DL.empty, DL.empty)
    simplifyM (Just mb) = do
        (b', st3, cdecl, vdecl) <- simplifyExpChoice (annL RInt mb)
        return (Just (unTypL b'), st3, cdecl, vdecl)

simplifyStmt l (Nop a) = return (DL.singleton (L l (Nop a)), DL.empty, DL.empty)

{-
Note:
Declarations inside sequences can depend on the iteration variable. Moreover,
declarations can also depend on constants which depend themselves on the 
iteration variable.
This way, we have to compute the transitive closure of dependencies starting
in the iteration variable. After having the set of all constants, we can also
determine the list of all variables which have dependencies. 
All variables whose type dependends on the iteration variable must remain
inside the body of the sequence. All other can be removed to outside the body
and shared between all iterations.
-}
declaredConsts
    :: [ConstDef]
    -> [Var]
declaredConsts = catMaybes . map (constDecl . unLoc)
    where
    constDecl (CDecl (ConstD (L _ cd) _ _)) = Just cd
    constDecl _ = Nothing 

innerConsts
    :: Var -> [ConstDef]
    -> ([ConstDef], [ConstDef])
innerConsts i cdecls = let
        -- The base case are the constants that depend on the index variable
        (base, rest) = Lst.partition (Set.member i . fvs) cdecls
    in fixpoint rest base
    where
    fixpoint rest [] = ([], rest)
    fixpoint rest base = let
            (base' , rest' ) = innerVars (declaredConsts base) rest
            (base'', rest'') = fixpoint rest' base'
        in (base ++ base'', rest'')

innerVars
    :: [Var] -> [VarDeclaration]
    -> ([VarDeclaration], [VarDeclaration])
innerVars consts vdecls = Lst.partition (mbr . fvs) vdecls
    where
    mbr vs = any (\c -> Set.member c vs) consts

{-
Note 3:
If we face a condition on a while statement which is not in the simplified form,
 we must simplify it to basic operations. However, unlike if statements, it 
is not enough to add them before the beginning of the cycle. We have also to 
add this to the end of the body of the cycle, because the condition has to be 
calculated in every iteration. Otherwise, we just have an infinit loop whenever 
the condition is true for the first values.

For instance:
@
 while (3 * i + j <  i * j) {
   ...
   i := i + 1;
   j := j + 1; 
}
@
must be simplified to:
@
 t0 := 3 * i;
 t1 := t0 + j;
 t2 := i * j;
 cond := t1 < t2;
 while (cond) {
   ...
   i := i + 1;
   j := j + 1;
   t0 := 3 * i;
   t1 := t0 + j;
   t2 := i * j;
   cond := t1 < t2;
}
@
-}

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

simplifyLValue 
    :: CaoMonad m 
    => LVal Var 
    -> m ( LVal Var
         , DL.DList SS
         , DL.DList ConstDef
         , DL.DList VarDeclaration)
simplifyLValue (LVVar (L l v)) = do
    (v', cdecl) <- simplifyVar v
    return (LVVar (L l v'), DL.empty, cdecl, DL.empty)
simplifyLValue (LVStruct lv fld) = do
    (lv', stmts, cdecl, vdecl) <- simplifyLValue lv
    -- XXX: type annotation of fld
    return (LVStruct lv' fld, stmts, cdecl, vdecl)
simplifyLValue (LVCont ty lv p) = do
    (ty', cdeclt) <- simplifyType ty
    (lv', stmts1, cdecl1, vdecl1) <- simplifyLValue lv
    (p' , stmts2, cdecl2, vdecl2) <- simplifyPat p
    return ( LVCont ty' lv' p'
           , stmts1 `DL.append` stmts2
           , cdeclt `DL.append` cdecl1 `DL.append` cdecl2
           , vdecl1 `DL.append` vdecl2)

-- Assignments -----------------------------------------------------------------

simplifyAssignment 
    :: CaoMonad m 
    => SrcLoc -> LVal Var -> TLExpr Var 
    -> m ( DL.DList (LStmt Var)
         , DL.DList ConstDef
         , DL.DList VarDeclaration)
simplifyAssignment loc lv e = do
    (lv', stmts1, cdecl1, vdecl1) <- simplifyLValue lv
    (e',  stmts2, cdecl2, vdecl2) <- 
-- When we have a simple left variable, we just have to simplify the assigned 
-- expression, and add a new assignment in the end. The use of 'simplifyExp'
-- guarantees that, for instance, 3 + v, is not further simplified.
        if' (isSimpleLVal lv') simplifyExp simplifyExpChoice e
    return ( stmts1 `DL.append` stmts2 `DL.snoc` L loc (Assign (lv':[]) (e':[]))
           , cdecl1 `DL.append` cdecl2
           , vdecl1 `DL.append` vdecl2)

{-

Note 6:
The simplification of parallel assignments is trickier because of
its semantics. The assigned value is always the value before the 
assignment. Thus, the following example:
@ a, b := b, a; @
is, in fact, the swap of the values between variables 'a' and 'b'.
This has to be expanded to:
@
 t0 := b;
 t1 := a;
 b  := t1;
 a  := t0;
@
to maintain the semantics. However, the code is more complex
and more variables are introduced.
-}

-- Precondition: |lvs| > 1, |exps| > 1, |lvs| = |exps|
simplifyMultipleAssignment 
    :: CaoMonad m 
    => SrcLoc -> [LVal Var] -> [TLExpr Var] 
    -> m (DL.DList (LStmt Var), DL.DList ConstDef, DL.DList VarDeclaration)
simplifyMultipleAssignment _ [] [] = return (DL.empty, DL.empty, DL.empty)
simplifyMultipleAssignment loc (lv:lvs) (e:exps) = do
    (lv', vdecl) <- newLVar (typeOf lv)
    -- "Frozzing" the values
    (stmt1, cdecl1, vdecl1) <- simplifyAssignment loc lv' e
    -- Handling the rest of the variables
    (stmt2, cdecl2, vdecl2) <- simplifyMultipleAssignment loc lvs exps
    -- Assigning the values
    (stmt3, cdecl3, vdecl3) <- simplifyAssignment loc lv (toExp lv')
    return ( stmt1  `DL.append` stmt2  `DL.append` stmt3
           , cdecl1 `DL.append` cdecl2 `DL.append` cdecl3
           , vdecl   `DL.cons` (vdecl1 `DL.append` vdecl2 `DL.append` vdecl3))
simplifyMultipleAssignment _ _ _ = caoError defSrcLoc $ mkUnknownErr
        "<Language.CAO.Transformation.Simplify>.\
            \<simplifyMultipleAssignment>: not expected case"


-- Precondition: |lvs| > 1 |e| = 1, e is a function call
simplifyTupleAssignment 
    :: CaoMonad m 
    => SrcLoc -> [LVal Var] -> TLExpr Var 
    -> m (DL.DList (LStmt Var), DL.DList ConstDef, DL.DList VarDeclaration)
-- We need a special case for simultaneous casts of function results, since 
-- this cannot be handled by 'simplifyExp'.
simplifyTupleAssignment loc lvs expr =
    case expr of
        L lc (TyE _ (Cast b tds@(_:_:_) ex@(L _ (TyE _ (FunCall _ _))))) -> do
            (ex' , stmts , cdecl1, vdecl1) <- simplifyExp ex
            (lvs', assign, cdecl2, vdecl2) <- 
                concatMapAndUnzip4MD (auxCast lc b) $
                    zip3 tds lvs (fromTuple $ typeOf ex)
            return ( stmts  `DL.append` (genLoc (Assign (DL.toList lvs') [ex']) 
                                        `DL.cons` assign)
                   , cdecl1 `DL.append` cdecl2
                   , vdecl1 `DL.append` vdecl2)
        e -> do
            (e'  , stmts , cdecl1, vdecl1) <- simplifyExp e
            (lvs', assign, cdecl2, vdecl2) <- concatMapAndUnzip4MD auxLv lvs
            return ( stmts  `DL.append` (genLoc (Assign (DL.toList lvs') [e'])
                                        `DL.cons` assign)
                   , cdecl1 `DL.append` cdecl2
                   , vdecl1 `DL.append` vdecl2)
    where
    auxCast lc b (td, lv, te) = let
            tlv = typeOf lv
            -- TODO: The sintactic equalify is too weak
            -- The typechecker could provide an annotation
        in if tlv == te
            then auxLv lv 
            else do
                (lv', ldecl) <- newLVar te
                (assign, cdecl, vdecl) <- simplifyAssignment loc lv $
                    annL tlv $ L lc $ Cast b [td] (toExp lv')
                return (DL.singleton lv', assign, cdecl, ldecl `DL.cons` vdecl)

    auxLv lv = if isSimpleLVal lv
        then return (DL.singleton lv, DL.empty, DL.empty, DL.empty)
        else do
             (lv', ldecl) <- newLVar $ typeOf lv
             (assign, cdecl, vdecl) <- simplifyAssignment loc lv (toExp lv')
             return (DL.singleton lv', assign, cdecl, ldecl `DL.cons` vdecl)
    
-- ConstDef --------------------------------------------------------------------

simplifyConstDeclaration 
    :: CaoMonad m 
    => ConstDecl Var -> Bool 
    -> m (ConstDecl Var, DL.DList ConstDef)
simplifyConstDeclaration (ConstD (L l n) b ce) _ =
    case ce of
        ConstInit _ -> do
            let Just e = indConst n
            (index, cdecl) <- simplifyIndexChoice e
            let n' = setIndConst index n
            return (ConstD (L l n') b (ConstInit (ind2Expr index)), cdecl)
        _ -> return (ConstD (L l n) b None, DL.empty)
simplifyConstDeclaration _ _ = internalError 
    "simplifyConstDeclaration" "Not expected multiple constant declarations"

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

{-
Note 8:
The declaration of variables, may include an optional definition, that should
not appear in the simplified form.

The type of variables must also be simplified so that later usage may benifit
of an already simplified type.


In global variables which are containers, the declaration, is like an 
assignment, and the order has to be preserved: 
- declaration of auxiliary variables
- simplification of values
- declaration of an auxiliary container of the same type, initialized with
  the simplified values
- assignment of the auxiliary container to the global container

For instance, in the declaration:
@ def v3 : vector[3] of register int := { a, b, v1[a] }; @

we have this generated code in the body of the init procedure:
@
 def c_t53 : register int;
 c_b := 3;
 c_t53 := c_v1[c_a];
 def c_t54 : vector[3] of register int := {c_a, c_b, c_t53};
 c_v3 := c_t54;
@

Putting the declaration next to the other declarations, we would obtain:
@
 def c_t53 : register int;
 def c_t54 : vector[3] of register int := {c_a, c_b, c_t53};
 c_b := 3;
 c_t53 := c_v1[c_a];
 c_v3 := c_t54;
@
This uses variables before their definition, namely 'c_b' and 'c_t53'.
-}
-- Global Variables
-- XXX: b -> type2TypeDecl?
simplifyVarDeclaration 
    :: CaoMonad m
    => VarDecl Var 
    -> m ( [VarDecl Var]
         , [SS]
         , [ConstDef]
         , [VarDeclaration]
         )
simplifyVarDeclaration (VarD (L l n) d Nothing) = do
    (n', cdecl) <- simplifyVar n
    return ([VarD (L l n') d Nothing], [], DL.toList cdecl, [])
simplifyVarDeclaration (VarD (L l x) b (Just e)) = do
    (x', cdecl1) <- simplifyVar x
    (e', ss, cdecl2, vdecl) <- simplifyExpChoice e
    let assign = genLoc $ Assign [LVVar (L l x')] [e']
    return ( VarD (L l x') b Nothing : []
           , DL.toList $ ss `DL.snoc` assign
           , DL.toList $ cdecl1 `DL.append` cdecl2
           , DL.toList vdecl)
simplifyVarDeclaration (MultiD xs b) = 
    concatMapAndUnzip4M (\ x -> simplifyVarDeclaration (VarD x b Nothing)) xs

simplifyVarDeclaration (ContD (L l lx) b es) = do
    let ty = varType lx
    (ty', cdecl1) <- simplifyType ty
    let lx' = L l $ setType ty' lx
    (es', ss, cdecl2, vdecl) <- simplifyExps es
    tv <- freshVar Local ty'
    return ( VarD lx' b Nothing : []
           , DL.toList $ ss 
                `DL.snoc` 
                -- Declaration used as assignment [See note 8]
                (genLoc $ VDecl  $ ContD (genLoc tv) b (DL.toList es'))
                `DL.snoc`
                (genLoc $ Assign [LVVar lx'] [genLoc $ annTyE ty' $ Var tv ])
           , DL.toList $ cdecl1 `DL.append` cdecl2
           , DL.toList vdecl
           )

-- Local Variables
simplifyLocalVarDeclaration 
    :: CaoMonad m
    => SrcLoc -> VarDecl Var 
    -> m (DL.DList (LStmt Var), DL.DList ConstDef, DL.DList VarDeclaration)
simplifyLocalVarDeclaration loc (VarD (L l n) _ Nothing) = do
    (n', cdecl) <- simplifyVar n
    --- XXX: reTypVar??
    reTypVar n'
    return (DL.singleton $ L loc $ VDecl $ VarD (L l n') (type2TyDecl (varType n')) Nothing
           , cdecl
           , DL.empty)
simplifyLocalVarDeclaration loc (VarD (L l x) _ (Just e)) = do
    (x', cdecl) <- simplifyVar x
    (e', ss, cdecl2, vdecls) <- simplifyExpChoice e
    let assign = genLoc $ Assign [LVVar (L l x')] [e']
    reTypVar x'
    return (L loc (VDecl $ VarD (L l x') (type2TyDecl (varType x')) Nothing)
                    `DL.cons` (ss `DL.snoc` assign)
           , cdecl `DL.append` cdecl2
           , vdecls)
simplifyLocalVarDeclaration loc (MultiD xs b) = do
    concatMapAndUnzip3MD 
        (\ x -> simplifyLocalVarDeclaration loc (VarD x b Nothing)) xs
simplifyLocalVarDeclaration loc (ContD (L l x) _ es) = do
    (x', cdecl) <- simplifyVar x
    (es', ss, cdecl2, vdecls) <- simplifyExps es
    reTypVar x'
    return ( ss `DL.snoc` 
                (L loc $ VDecl $ ContD (L l x') (type2TyDecl (varType x')) (DL.toList es'))
           , cdecl `DL.append` cdecl2
           , vdecls)

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

simplifyExps 
    :: CaoMonad m 
    => [TLExpr Var] 
    -> m ( DL.DList (TLExpr Var)
         , DL.DList SS
         , DL.DList ConstDef
         , DL.DList VarDeclaration
         )
simplifyExps = fold4M 
    simplifyExpChoice 
    DL.cons DL.append DL.append DL.append 
    (DL.empty, DL.empty, DL.empty, DL.empty)

{-
Note 4:
There are two simplification functions to expressions, with a little different
behavior: (the naming is not the better one)
- simplifyExp: Expressions are only simplified if they are operations on other
    operations. This means that variables and constants as operands are left
    as they are.
- simplifyExpChoice: Only constants and variables are left as they are; 
    all other expressions are assigned to a new variable.
-}
simplifyExp 
    :: CaoMonad m 
    => TLExpr Var 
    -> m (TLExpr Var, DL.DList SS, DL.DList ConstDef, DL.DList VarDeclaration)
simplifyExp (L l e) = do
    (e', as, cdecl, vdecl) <- simplExp e
    return (L l e', as, cdecl, vdecl)

-- XXX: should the type annotation be modified using simplifyType?
simplExp 
    :: CaoMonad m 
    => TExpr Var 
    -> m (TExpr Var, DL.DList SS, DL.DList ConstDef, DL.DList VarDeclaration)
simplExp (TyE t l@(Lit _)) = do
    (t', cdecl) <- simplifyType t
    return (TyE t' l, DL.empty, cdecl, DL.empty)
simplExp (TyE _ (Var   v)) = do
    (v', cdecl) <- simplifyTVar v
    return (v', DL.empty, cdecl, DL.empty)
simplExp (TyE t (FunCall f es)) = do 
-- XXX: annotation on f?
    (t', cdecl) <- simplifyType t
    (es', stmts, cdecls, vdecls) <- simplifyExps es
    return ( TyE t' (FunCall f (DL.toList es'))
           , stmts
           , cdecl `DL.append` cdecls
           , vdecls)
simplExp (TyE t (StructProj ea n)) = do
-- XXX: annotation on n
    (t', cdecl) <- simplifyType t
    (ea', stmts, cdecls, vdecls) <- simplifyExpChoice ea
    return (TyE t' (StructProj ea' n), stmts, cdecl `DL.append` cdecls, vdecls)
simplExp (TyE t (UnaryOp op e)) = do
    (t', cdecl) <- simplifyType t
    (e', ss, cdecls, vdecls) <- simplifyExpChoice e
    return (TyE t' (UnaryOp op e'), ss, cdecl `DL.append` cdecls, vdecls)
simplExp (TyE t (BinaryOp op l r)) = do
    (t', cdecl) <- simplifyType t
    ((l',r'), ss, cdecls, vdecls) <- simplifyBinaryExp l r
    return (TyE t' (BinaryOp op l' r'), ss, cdecl `DL.append` cdecls, vdecls)
simplExp (TyE ty (Access e p)) = do
    (ty', cdecl) <- simplifyType ty
    (e', ss1, cdecls1, vdecls1) <- simplifyExpChoice e
    (p', ss2, cdecls2, vdecls2) <- simplifyPat p
    return (TyE ty' (Access e' p')
           , ss1     `DL.append` ss2
           , cdecl   `DL.append` cdecls1 `DL.append` cdecls2
           , vdecls1 `DL.append` vdecls2)
-- XXX: update type declaration
simplExp (TyE ty (Cast b td e)) = do
    (ty', cdecl) <- simplifyType ty
    (e', stmts, cdecls, vdecl) <- simplifyExpChoice e
    return ( TyE ty' (Cast b td e')
           , stmts
           , cdecl `DL.append` cdecls
           , vdecl)

-- Simplifies both operands of a binary expression
simplifyBinaryExp :: CaoMonad m => TLExpr Var -> TLExpr Var
                  -> m ((TLExpr Var, TLExpr Var)
                       , DL.DList SS
                       , DL.DList ConstDef
                       , DL.DList VarDeclaration)
simplifyBinaryExp l r = do
    (l', stmts1, index1, decl1) <- simplifyExpChoice l
    (r', stmts2, index2, decl2) <- simplifyExpChoice r
    return ((l', r')
           , stmts1 `DL.append` stmts2
           , index1 `DL.append` index2
           , decl1  `DL.append` decl2)

simplifyExpChoice 
    :: CaoMonad m
    => TLExpr Var 
    -> m (TLExpr Var, DL.DList SS, DL.DList ConstDef, DL.DList VarDeclaration)
-- Literals: there is no need to introduce a new variable
simplifyExpChoice (L loc (TyE t l@(Lit _))) = do
    (t', cdecl) <- simplifyType t
    return (L loc (TyE t' l), DL.empty, cdecl, DL.empty)
-- Variables: there is no need to introduce a new variable. Type annotations
-- are updated.
simplifyExpChoice (L l (TyE _ (Var v))) = do
    (v', cdecl) <- simplifyTVar v
    return (L l v', DL.empty, cdecl, DL.empty)
simplifyExpChoice e = do
    (e', stmts, cdecl, decl1) <- simplifyExp e
    (ve, assign, decl2) <- assignToNewVar e'
    return (ve, stmts `DL.snoc` assign, cdecl, decl1 `DL.snoc` decl2)


--------------------------------------------------------------------------------
-- Accesses: just boilerplate

simplifyPat 
    :: CaoMonad m => APat Var
    -> m (APat Var, DL.DList SS, DL.DList ConstDef, DL.DList VarDeclaration)
simplifyPat (VectP r) = do
    (r', ss, cdecls, vdecls) <- simplifyRowPat r
    return (VectP r', ss, cdecls, vdecls)
simplifyPat (MatP r c) = do
    (r', ss1, cdecls1, vdecls1) <- simplifyRowPat r
    (c', ss2, cdecls2, vdecls2) <- simplifyRowPat c
    return ( MatP r' c'
           , ss1     `DL.append` ss2
           , cdecls1 `DL.append` cdecls2
           , vdecls1 `DL.append` vdecls2)

simplifyRowPat 
    :: CaoMonad m 
    => RowAPat Var 
    -> m (RowAPat Var, DL.DList SS, DL.DList ConstDef, DL.DList VarDeclaration)
simplifyRowPat (CElem e) = do
    (e', ss, cdecls, vdecls) <- simplifyExpChoice e
    return (CElem e', ss, cdecls, vdecls)
simplifyRowPat (CRange e1 e2) = do
    (e1', ss1, cdecls1, vdecls1) <- simplifyExpChoice e1
    (e2', ss2, cdecls2, vdecls2) <- simplifyExpChoice e2
    return ( CRange e1' e2'
           , ss1     `DL.append` ss2
           , cdecls1 `DL.append` cdecls2
           , vdecls1 `DL.append` vdecls2)

--------------------------------------------------------------------------------
-- Types

simplifyVar 
    :: CaoMonad m 
    => Var
    -> m (Var, DL.DList ConstDef)
simplifyVar v = do 
    mv <- lookupReTypVar v
    case mv of
        Nothing -> do
            (t, cdecl) <- simplifyType $ varType v
            return (setType t v, cdecl)
        Just v' -> return (v', DL.empty)
    
simplifyTVar
    :: CaoMonad m 
    => Var
    -> m (TExpr Var, DL.DList ConstDef)
simplifyTVar v = do
    (v', cdecl) <- simplifyVar v
    return (annTyE (varType v') $ Var v', cdecl)


simplifyType
    :: CaoMonad m 
    => Type Var
    -> m (Type Var, DL.DList ConstDef)
simplifyType (Tuple tlst) = do
    (tlst', cdecls) <- fold2M' simplifyType (flip DL.snoc) DL.append (DL.empty, DL.empty) tlst
    return (Tuple $ DL.toList tlst', cdecls)
simplifyType (Bits s n) = do
    (n', cdecls) <- simplifyIndexChoice n
    return (Bits s n', cdecls)
simplifyType (Vector n t) = do
    (n', cdecls1) <- simplifyIndexChoice n
    (t', cdecls2) <- simplifyType t
    return (Vector n' t', cdecls1 `DL.append` cdecls2)
simplifyType (Matrix n m t) = do
    (n', cdecls1) <- simplifyIndexChoice n
    (m', cdecls2) <- simplifyIndexChoice m
    (t', cdecls3) <- simplifyType t
    return (Matrix n' m' t', cdecls1 `DL.append` cdecls2 `DL.append` cdecls3)

simplifyType (Mod Nothing Nothing (Pol [Mon (CoefI m) EZero])) = do
    (m', cdecls) <- simplifyIndexChoice m
    return (Mod Nothing Nothing (Pol [Mon (CoefI m') EZero]), cdecls)
    
simplifyType t = return (t, DL.empty)

simplifyIndex 
    :: CaoMonad m 
    => IExpr Var 
    -> m (IExpr Var, DL.DList ConstDef)
simplifyIndex n@(IInt _) = return (n, DL.empty)
simplifyIndex v@(IInd _) = return (v, DL.empty)
simplifyIndex (IArith op e1 e2) = do
    (e1', stmts1) <- simplifyIndexChoice e1
    (e2', stmts2) <- simplifyIndexChoice e2
    return (IArith op e1' e2', stmts1 `DL.append` stmts2)
simplifyIndex (ISym e) = do
    (e', stmts)   <- simplifyIndexChoice e
    return (ISym e', stmts)
simplifyIndex (ISum slst) = simplifySum slst

-- This function takes a sum of terms and returns a tree of binary additions.
simplifySum 
    :: CaoMonad m 
    => [IExpr Var] -- List of terms
    -> m (IExpr Var, DL.DList ConstDef)
simplifySum []       = internalError "simplifySum" "Empty sum!"
-- When we have the sum of just one term, we can remove the sum
simplifySum [e]      = simplifyIndex e
-- The general base case has two terms
simplifySum [e1, e2] = do
    (e1', stmts1)   <- simplifyIndexChoice e1
    (e2', stmts2)   <- simplifyIndexChoice e2
    return (ISum [e1', e2'], stmts1 `DL.append` stmts2)
simplifySum (e:lest) = do
    (e', stmts1)    <- simplifyIndexChoice e
    (lest', stmts2) <- simplifySum lest
    (iv, cdecl)     <- newIndexDef lest'
    return (ISum [e', iv], stmts1 `DL.append` stmts2 `DL.snoc` cdecl)

simplifyIndexChoice 
    :: CaoMonad m 
    => IExpr Var 
    -> m (IExpr Var, DL.DList ConstDef)
simplifyIndexChoice n@(IInt _) = return (n, DL.empty)
simplifyIndexChoice v@(IInd _) = return (v, DL.empty)
simplifyIndexChoice e = do
    (e', cdecls) <- simplifyIndex e
    (iv, cdecl)  <- newIndexDef e'
    return (iv, cdecls `DL.snoc` cdecl)

--------------------------------------------------------------------------------
-- Variable generation

-- Generates a new local variable with the respective declaration.
newVariable :: CaoMonad m => Type Var -> m (Var, VarDeclaration)
newVariable typ = do
    tv <- freshVar Local typ
    let decl = genLoc $ VDecl $ VarD (genLoc tv) (type2TyDecl typ) Nothing
    return (tv, decl)

-- New left variable with the respective declaration
newLVar :: CaoMonad m => Type Var -> m (LVal Var, VarDeclaration)
newLVar = liftM (mapFst (LVVar . genLoc)) . newVariable 

-- Given an expression, returns a new variable with the same type, together with
-- its declaration and an assignment of the expression.
-- E.g.
-- Literal 3 of type int
-- assignToNewVar 3 -> (t0, def t0 : int, t0 := 3)
-- Sum of two integers
-- assignToNewVar (3+i) -> (t0, def t0: int, t0 := 3 + i)
-- Here, SS introduces an assignment
assignToNewVar :: CaoMonad m => TLExpr Var -> m (TLExpr Var, SS, VarDeclaration)
assignToNewVar e = do
    let ty = typeOf e
    (tv, decl) <- newVariable ty
    let assign = genLoc $ Assign [LVVar (genLoc tv)] [e]
    return (genLoc $ TyE ty $ Var tv, assign, decl)

{-
Note 7:
The first version of this functions was introducing ordinary variables to
simplify type expression. Altough this would generate valid C code (because
the static library does not know anything about index constants), this breaks
the correctness of the intermediate CAO code. Moreover, the optimization stage
loses this important meta-information.

To maintain the correctness of the intermediate CAO program, a new local index 
constant has to be declared and defined. Since, by definition, constants cannot
be assigned, their value has to be set during declaration (declaration and 
definition are simultaneous), and can only depend on other constants.
-}

newIndexDef :: CaoMonad m => IExpr Var -> m (IExpr Var, ConstDef)
newIndexDef e = do
    let ty = typeOf e
    tv <- freshIndex Local ty
    let decl = genLoc $ CDecl $ 
            ConstD (genLoc tv) (type2TyDecl ty) (ConstInit (ind2Expr e))
    return (IInd tv, decl)

--------------------------------------------------------------------------------
-- Auxiliary functions

toExp :: LVal Var -> TLExpr Var
toExp (LVVar (L l v)) = L l $ annTyE (varType v) $ Var v
toExp _         = error "<Language.CAO.Transformation.Simplify>.\
    \<toExp>: undefined case"

moduleName :: String
moduleName = "<Language.CAO.Transformation.Simplify>"

internalError :: String -> String -> a
internalError funcName msg = error $
    moduleName ++ ".<" ++ funcName ++ ">: " ++ msg