{-  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 :  Sequence unrolling.
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

In the expansion phase, sequential code replaces iteration in sequence 
statements. This phase is optional: the sequence statements can be later 
translated to iterative C code. Expansion of sequences is a trade-off between 
the memory used by the machine code and the execution time. Usually expanded 
code will be faster because there are no conditional jumps and some of the 
expressions can be partially evaluated. However, this may not be the case if 
the target machine architecture uses an intermediate cache memory that is not 
enough to hold all the code. In this situation, conditional jumps may be 
preferable to cache misses but this has to be determined experimentally.

A sequence statement is an iteration instruction where the bounds and the 
increment of the index (bound) variable are statically known. This means that 
we can compute during compilation the number of times that the sequence body is 
executed and the values that the index variable will take. To expand the 
sequence, its body is replicated by that number of times and the sequence index 
is replaced by its respective value. Although similar to traditional loop 
unrolling, this expansion has some subtleties:

* In nested sequences, for each value taken by the index variable of the outer 
sequence, there has to be a list of index variable values for the inner 
sequence.  This implies that the outer sequence has to be expanded before the 
inner sequence.

* Subsequent steps rely on type annotations to generate correctly typed code, 
thus type annotations in expanded code must be updated accordingly with 
expansion. Since CAO has a limited form of dependent types, the type of some 
expressions inside the sequence body are functions of the index variable.

-}

module Language.CAO.Transformation.Expand (
    expandSequences
) where

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

import Data.DList ( DList )
import qualified Data.DList as DL
import Data.Set ( Set )
import qualified Data.Set as Set

import Language.CAO.Common.Literal
import Language.CAO.Common.Monad
import Language.CAO.Common.Polynomial
import Language.CAO.Common.SrcLoc
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.Type

-- | This function expands the body sequence statements with known bounds.
--   If any limit (bounds) is not statically known, the sequence body
--   is not expanded.
expandSequences :: CaoMonad m => Prog Var -> m (Prog Var)
expandSequences (Prog defs _) = 
    liftM2 Prog (mapM (mapML go) defs) (return Nothing)
    where 
    -- Simple program traversal to handle with statements
    go :: CaoMonad m => Def Var -> m (Def Var)
    go (FunDef (Fun n args rt ss)) = 
        FunDef . Fun n args rt <$> concatMapM expandStmt ss
    go d = return d

-- Since a single sequence statement can be expanded to a block of statements, 
-- the resulting type is a list
-- We must also traverse statements which contain themselves blocks of 
-- statements.
expandStmt :: CaoMonad m => LStmt Var -> m [LStmt Var]
expandStmt s@(L _ (Seq _ _)) = seqCase s 
expandStmt (L l (Ite i t e)) =
    singleton . L l <$> liftM2 (Ite i) (concatMapM expandStmt t) 
                                       (mapMaybeM (concatMapM expandStmt) e)
expandStmt (L l (While c ss))    = 
    singleton . L l . While c <$> concatMapM expandStmt ss
expandStmt s = return [s]

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

-- Values that the bound variable will take during the sequence execution
seqRange :: Integer -> Integer -> Integer -> [Integer]
seqRange strt final dist = enumFromThenTo strt (strt + dist) final

seqCase :: CaoMonad m => 
    LStmt Var -> m [LStmt Var]
seqCase (L loc (Seq (SeqIter ivar estart eend eby rng) ss)) = do
    case (unLoc estart, unLoc eend) of
        -- The bounds are statically known
        (Lit (ILit estart'), Lit (ILit eend')) -> do
            let insts = seqRange estart' eend' (auxMBy eby)
                bvars  = bvs ss
            -- Expands the sequence:
            stmt <- expandSeq ss bvars ivar insts 
            -- Expands nested sequences:
            concatMapM expandStmt stmt 
        -- The bounds are not statically knonw, but inner sequences must
        -- be expanded
        _ -> singleton . L loc . Seq (SeqIter ivar estart eend eby rng) <$> 
                concatMapM expandStmt ss
    where
    auxMBy Nothing                      = 1
    auxMBy (Just (L _ (Lit (ILit by)))) = by
    auxMBy e                            = error $ show e
seqCase _ = error "<Language.CAO.Transformation.Expand>.\
    \<seqCase>: unexpected case"

expandSeq :: CaoMonad m => [LStmt Var] -> Set Var -> Var -> [Integer] -> m [LStmt Var]
expandSeq stmt bvars ivar ilst = liftM DL.toList $ foldM worker DL.empty ilst
    where
    worker :: CaoMonad m => DList (LStmt Var) -> Integer -> m (DList (LStmt Var))
    worker sstms i = do
        -- Gets a new unique identifier to each bound variable of the sequence
        -- XXX: do we need this?
        rbv <- mapM (\ x -> uniqId >>= \ i' -> return (x, i')) bvsSeq
        return $ sstms `DL.append` DL.fromList (renameStmt rbv i)

    -- XXX: is this definitions correct?
    bvsSeq :: [Var]
    bvsSeq = Set.toList bvars
    
    renameStmt :: [(Var, Int)] -> Integer -> [LStmt Var]
    renameStmt rbv i = map (sLStmt (ivar, IInt i) 
                           . (renamer $ retyp . renameBVs rbv)) 
                     $ subst (ivar, Lit $ ILit i) stmt
        where

        renamer :: (Var -> Var) -> LStmt Var -> LStmt Var
        renamer f = fmap (fmap f)

        -- Correcting type annotations, so that the index variable is replaced by
        -- its instantiation value
        retyp :: Var -> Var
        retyp v = setType (sType (ivar, IInt i) $ typeOf v) v

        renameBVs :: [(Var, Int)] -> Var -> Var
        renameBVs bvslst v = maybe v (flip setId v) (lookup v bvslst)

--------------------------------------------------------------------------------
-- More boilerplate...
-- This should be replaced by a generic transformation

sLStmt :: (Var, IExpr Var) -> LStmt Var -> LStmt Var
sLStmt s = fmap (sStmt s)

sStmt :: (Var, IExpr Var) -> Stmt Var -> Stmt Var
sStmt s (Assign lvals es) = Assign (map (sLVal s) lvals) (map (sTLExpr s) es)
sStmt s (FCallS f es) = FCallS f (map (sTLExpr s) es)
sStmt s (Ret es) = Ret (map (sTLExpr s) es)
sStmt s (Ite e stmts mst) = Ite (sTLExpr s e) (map (sLStmt s) stmts) (fmap (map (sLStmt s)) mst)
sStmt s (While e stmts) = While (sTLExpr s e) (map (sLStmt s) stmts)
sStmt s (Seq iter stmts) = Seq iter (map (sLStmt s) stmts)
sStmt _ s = s

sTLExpr :: (Var, IExpr Var) -> TLExpr Var -> TLExpr Var
sTLExpr s (L l (TyE t e)) = L l $ TyE (sType s t) (sExpr s e)

sExpr :: (Var, IExpr Var) -> Expr Var -> Expr Var
sExpr s (FunCall f es) = FunCall f (map (sTLExpr s) es)
sExpr s (StructProj e fld) = StructProj (sTLExpr s e) fld
sExpr s (UnaryOp op e) = UnaryOp op (sTLExpr s e)
sExpr s (BinaryOp op e1 e2) = BinaryOp op (sTLExpr s e1) (sTLExpr s e2)
sExpr s (Access e pat) = Access (sTLExpr s e) pat
sExpr s (Cast b d e) = Cast b d (sTLExpr s e)
sExpr _ e = e

sLVal :: (Var, IExpr Var) -> LVal Var -> LVal Var
sLVal s (LVVar (L l v)) = LVVar $ L l $ setType (sType s $ typeOf v) v
sLVal s (LVStruct lv fld) = LVStruct (sLVal s lv) fld
sLVal s (LVCont typ lv pat) = LVCont (sType s typ) (sLVal s lv) pat

sType :: (Var, IExpr Var) -> Type Var -> Type Var
sType s (Bits sg e) = Bits sg $ evalExpr (subst s e)
sType s (Mod Nothing Nothing (Pol [Mon (CoefI m) EZero])) =
    Mod Nothing Nothing (Pol [Mon (CoefI (evalExpr (subst s m))) EZero])
sType s (Vector e t) = Vector (evalExpr (subst s e)) (sType s t)
sType s (Matrix e1 e2 t) = Matrix (evalExpr (subst s e1)) (evalExpr (subst s e2)) (sType s t)
sType s (Tuple ts) = Tuple $ map (sType s) ts
sType _ t = t
-- XXX: This definition is incomplete and may have some problems with indexes and mods