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

{-
Module      :  $Header$
Description :  CAO static single assignment form.
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

CAO static single assignment form.
-}

module Language.CAO.Analysis.SSA
    ( toSSA
    , fromSSA
    ) where

import Control.Monad.State.Strict
import Data.Graph
import Data.List hiding (insert) 
import Data.Map (Map)
import qualified Data.Map as Map

import Language.CAO.Analysis.CFG
import Language.CAO.Analysis.Dominance
import Language.CAO.Analysis.PhiInsert
import qualified Language.CAO.Analysis.SsaBack as BT

import Language.CAO.Common.Monad
import Language.CAO.Common.Var
import Language.CAO.Common.SrcLoc
import Language.CAO.Common.State
import Language.CAO.Common.Utils ( mapFst, mapSnd, replaceAt )

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

import Language.CAO.Type


---- Renaming Variables --------------------------------------------------------

-- Local definition
newtype Stack a = Stack [a]

push :: a -> Stack a -> Stack a
push a (Stack s) = Stack (a : s)

pop :: Stack a -> (a, Stack a)
pop (Stack []) = error "Pop: Empty stack."
pop (Stack (h:t)) = (h, Stack t)

top :: Stack a -> Maybe a
top (Stack []) = Nothing
top (Stack s) = Just $ head s

emptyStack :: Stack a
emptyStack = Stack []

-- count: C(*) map from variable identifiers (v) to a counter telling
--      how many assignments to v have been processed
-- stacks: S(*) map from variable identifiers (v) 
--      to stacks of integers representing ...
data RenameState =
    RState { count  :: !Int
           , stacks :: !(Map Var (Stack Int))
           }

-- Top-level
emptyRState :: Int -> RenameState
emptyRState vuniq = RState vuniq Map.empty

popM :: Var -> State RenameState ()
popM str = modify $ \st -> snd $ pop_a st str
    where
-- Local definition
    pop_a :: RenameState -> Var -> (Int,RenameState)
    pop_a st a = 
        let stacks_   = stacks st
            st_a      = Map.findWithDefault emptyStack a stacks_
            (x,st_a') = pop st_a
            st'       = Map.insert a st_a' stacks_
        in (x,st { stacks = st'})

pushM :: Var -> Int -> State RenameState ()
pushM str i = modify $ \st -> push_i_a st
    where
-- push i onto stack[a]
    push_i_a :: RenameState -> RenameState
    push_i_a st = 
        let stacks_ = stacks st
            st_a  = Map.findWithDefault emptyStack str stacks_
            st_a' = push i st_a
            st'   = Map.insert str st_a' stacks_ -- Use update in place ???
        in st { stacks = st' }

countM :: State RenameState Int
countM = do
    st <- get
    let c = count st
    put $ st { count = c + 1 }
    return c


-- Local definition
-- top stack[a]
top_a :: RenameState -> Var -> Maybe Int    
top_a st a = top $ Map.findWithDefault emptyStack a $ stacks st


--------------------------------------------------------------------
updateBlock :: NodeId -> CaoCFG -> [LStmt Var] -> CaoCFG
updateBlock nid cfg nstmts = 
    cfg { blocks = Map.adjust (mapFst (const nstmts)) nid (blocks cfg) }

blockById :: NodeId -> CaoCFG -> [LStmt Var]
blockById nid m = fst $ blocks m Map.! nid

            
------------------------------------------------------------------
-- TODO: FIX -> monadic uniq indentifiers for renaming!!!
renameVars :: CaoMonad m => Map Vertex Vertex -> CaoCFG -> [Var] -> m CaoCFG
renameVars domTree cfg vs = do
    u <- uniqId
    let initSt        = foldl' aux (emptyRState u) vs
        (cfg', st')   = runState (rename (invertMap domTree) cfg entryNode) initSt
    st <- get
    put st { lastVar = count st' + 1 }
    return cfg'
    where 
    aux :: RenameState -> Var -> RenameState
    aux st v = st { stacks = Map.insert v emptyStack (stacks st) }
 
--------------------------------------------------------------------

rename :: Map Vertex [Vertex] -> CaoCFG -> NodeId -> State RenameState CaoCFG
rename domTree cfg nid = do
    (oldLHS, cfg') <- blockAssignments cfg nid
    cfg''          <- foldM (phiFunctions nid) cfg' $ successors' nid cfg'
    cfg'''         <- foldM (rename domTree) cfg''  $ children domTree nid
    mapM_ popM oldLHS
    return cfg'''

-- First loop ------------------------------------------------------------------
blockAssignments :: CaoCFG -> NodeId -> State RenameState ([Var], CaoCFG)
blockAssignments cfg nid = do
    let stmtBlock = blockById nid cfg 
    (oldLHS, stmtBlock') <- renameStatements stmtBlock
    let cfg' = updateBlock nid cfg stmtBlock'
    return (oldLHS, cfg')

-- There is a similar function in Simplify module - fuse
--variablesLHS :: [LStmt Var] -> [Var]
--variablesLHS = Set.toList . lvalNames 

---- The CFG has empty lists of statements, thus we do not need recursion
---- The statements must be traversed in order
---- The RHS must be processed before the LHS
renameStatements :: [LStmt Var] -> State RenameState ([Var],[LStmt Var])
renameStatements = doMap
    where
    doMap xs = mapM aux xs >>= \lst ->
        let (a,b) = unzip lst in return (concat a, b)
    aux :: LStmt Var -> State RenameState ([Var], LStmt Var)
    aux s = case unLoc s of
        Assign lvs f@[unLoc -> unTyp -> FunCall n _] | isPhiFun (unLoc n) -> do
            lvs' <- mapM renameLVal lvs
            return ([],L (getLoc s) $ Assign lvs' f)
        Assign lvs rhs -> do
            rhs' <- mapM renameVar rhs
            lvs' <- mapM renameLVal lvs
            return (map lvname lvs, L (getLoc s) $ Assign lvs' rhs')
        FCallS fname exs -> do
            exs' <- mapM renameVar exs
            return ([],L (getLoc s) $ FCallS fname exs')
        Ret exs -> do
            exs' <- mapM renameVar exs
            return ([],L (getLoc s) $ Ret exs')
        Ite i t e -> do
            i' <- renameVar i
            return ([],L (getLoc s) $ Ite i' t e)
        While i ss -> do
            i' <- renameVar i
            return ([],L (getLoc s) $ While i' ss)

        Seq (SeqIter ivar ilow ihigh Nothing rng) stmts -> do
            ilow' <- renameVar' ilow
            ihigh' <- renameVar' ihigh
            return ([], L (getLoc s) $
                Seq (SeqIter ivar ilow' ihigh' Nothing rng) stmts)
        Seq (SeqIter ivar ilow ihigh (Just iby) rng) stmts -> do
            ilow' <- renameVar' ilow
            ihigh' <- renameVar' ihigh
            iby' <- renameVar' iby
            return ([], L (getLoc s) $
                Seq (SeqIter ivar ilow' ihigh' (Just iby') rng) stmts)

        Nop t -> return ([], L (getLoc s) $ Nop t)
        _         -> error "<Language.CAO.Analysis.SSA>.<renameStatements>:\
            \ FIXME! Add cases VDecl"
        -- FIXME: Not expecting sequences

    renameVar :: TLExpr Var -> State RenameState (TLExpr Var)
    renameVar (L l (TyE t e)) = get >>= \st ->
        return $ L l $ TyE t $ fmap (suffixVar st) e

    renameVar' :: LExpr Var -> State RenameState (LExpr Var)
    renameVar' (L l e) = get >>= \st ->
        return $ L l $ fmap (suffixVar st) e

    renameLVal :: LVal Var -> State RenameState (LVal Var)
    renameLVal lv = case lv of
        LVVar v -> do
            i <- countM
            pushM (unLoc v) i
            return $ LVVar $ fmap (setId i) v
        _ -> error "<SSA>.<renameLVal>: unexpected case" -- return lv


suffixVar :: RenameState -> Var -> Var
suffixVar st x
  | Just vid <- top_a st x = setId vid x 
  | otherwise              = x
--
--
-- Second loop -----------------------------------------------------------------
successors' :: NodeId -> CaoCFG -> [NodeId]
successors' nid = snd . (Map.! nid) . blocks

-- WhichPred(Y, X), Y in Succ(X)
whichPredecessor :: NodeId -> NodeId -> CaoCFG -> Int
whichPredecessor nsucc nid = aux 0 . Map.assocs . blocks
    where
    aux _ [] = error "<SSA>.<whichPredecessor>: empty list"
    aux n ((k, (_, succs)) : xs)
        | k == nid = n
        | nsucc `elem` succs = aux (n + 1) xs
        | otherwise = aux n xs

phiFunctions :: NodeId -> CaoCFG -> NodeId -> State RenameState CaoCFG
phiFunctions nid cfg nsucc = do
    st <- get
    let j = whichPredecessor nsucc nid cfg
    return $ updateBlock nsucc cfg $ 
             renamePhiFunc st j $ 
             blockById nsucc cfg 

renamePhiFunc :: RenameState -> Int -> [LStmt Var] -> [LStmt Var]
renamePhiFunc st j = map aux
    where
    aux :: LStmt Var -> LStmt Var
    aux s = case unLoc s of
        Assign lvs [L loc (TyE t (FunCall fname exps))] 
          | isPhiFun (unLoc fname) -> 
            L (getLoc s) $ Assign lvs [L loc (TyE t (FunCall fname $ jElem st j exps))]
        _ -> s
    jElem :: RenameState -> Int -> [TLExpr Var] -> [TLExpr Var]
    jElem st' j' exps = let
            L l (TyE t (Var v)) = exps !! j'
            v' = suffixVar st' v
        in replaceAt j' (L l $ TyE t $ Var v') exps

children :: Map Vertex [Vertex] -> NodeId -> [NodeId]
children = (Map.!)
--------------------------------------------------------------------------------
-- Store and Load --------------------------------------------------------------
--------------------------------------------------------------------------------

-- A function that takes in each node of the CFG, and removes every store and load into arrays and structures,
-- transforming them into function calls to 
-- store(vector/array_name,index/field_name,new value) and 
-- load(vector/array_name,index/field_name)

--------------------------------------------------------------------------------
-- Replaces writes/reads on global variables with procedure calls.
-- Writes and reads of structured types are replaced by function calls.

introLoadStore :: CaoCFG -> CaoCFG
introLoadStore cfg = cfg { blocks = blocks' }
    where
    wvars   = getWVars cfg
    lb      = loadBlock wvars
    sb      = storeBlock wvars
    blocks' = storeOnExit sb $ loadOnEntry lb $ Map.map (renameBlock lb sb) (blocks cfg)

--------------------------------------------------------------------------------
-- Written global variables in a CFG

getWVars :: CaoCFG -> [Var]
getWVars cfg = 
    case unLoc (definition cfg) of
        FunDef (Fun fn _ _ _) ->
            case varType (unLoc fn) of
                FuncSig _ _ (Proc wvars) -> wvars
                _ -> []
        _ -> []

--------------------------------------------------------------------------------
-- Block of global variable load and store statements

loadBlock :: [Var] -> [LStmt Var]
loadBlock = map aux
    where
    aux :: Var -> LStmt Var
    aux v = genLoc $ Assign [lv] [f]
        where
        f  = genLoc $ annTyE (varType v) $ FunCall lg []
        lg = genLoc $ mkLoadGlobal (varName v)
        lv = LVVar $ genLoc v


storeBlock :: [Var] -> [LStmt Var]
storeBlock = map $ \ v ->
    genLoc $ FCallS (mkStoreGlobal $ varName v) [genLoc $ annTyE (varType v) $ Var v]

--------------------------------------------------------------------------------
-- Loads written global variables on entry

loadOnEntry :: BasicBlock -> LocalGraph -> LocalGraph
loadOnEntry loadBlk blks = let
        i = head $ snd $ blks Map.! entryNode
    in Map.adjust (mapFst (loadBlk ++)) i blks

-- Stores written global variables on exit

storeOnExit :: BasicBlock -> LocalGraph -> LocalGraph
storeOnExit storeBlk = Map.adjust (mapFst (++ storeBlk)) exitNode

--------------------------------------------------------------------------------
-- Adds calls to store and retrieve global variables before and after function calls.
-- Replaces write/read to structured types by store/load function calls

renameBlock :: BasicBlock -> BasicBlock -> (BasicBlock, Connections) -> (BasicBlock, Connections)
renameBlock lb sb = mapFst (concatMap aux)
    where
    aux :: LStmt Var -> BasicBlock
    aux stmt = case unLoc stmt of
        -- Function Call
        FCallS _ _ -> sb ++ stmt : lb
        Assign _ [unLoc -> unTyp -> FunCall _ _] -> sb ++ stmt : lb
        -- Store
        Assign lv r            -> [ storeLoad lv r ]
        _ -> [stmt]

storeLoad :: [LVal Var] -> [TLExpr Var] -> LStmt Var
storeLoad lv exps = case head lv of
    LVVar _ -> genLoc $ Assign lv (map load exps)
    lv' -> let (lv'', lpath) = extractLVal lv'
        in storeCall lv'' lpath exps

load :: TLExpr Var -> TLExpr Var
load (L l (TyE t e)) = L l $ TyE t $ load_ e

load_ :: Expr Var -> Expr Var
load_ (StructProj s f)
    = FunCall (genLoc loadStruct) [s, genLoc $ annTyE (varType f) $ Var f]
load_ (Access c (VectP (CElem i)))
    = FunCall (genLoc loadVar) [c,i]
load_ (Access c (VectP (CRange i j)))
    = FunCall (genLoc loadVarRng) [c,i,j]
load_ (Access c (MatP (CElem i) (CElem j)))
    = FunCall (genLoc loadMatrix) [c,i,j]
load_ (Access c (MatP (CRange i j) (CRange k l)))
    = FunCall (genLoc loadMatrixRng) [c,i,j,k,l]
load_ (Access c (MatP (CRange i j) (CElem k)))
    = FunCall (genLoc loadMatrixRowRng) [c,i,j,k]
load_ (Access c (MatP (CElem i) (CRange j k)))
    = FunCall (genLoc loadMatrixColRng) [c,i,j,k]
load_ e
    = e

storeCall :: Var -> [TLExpr Var] -> [TLExpr Var] -> LStmt Var
storeCall lv index values = let 
        lv'  = LVVar $ genLoc lv
        dest = genLoc $ annTyE (varType lv) $ Var lv
        funC = genLoc $ annTyE (varType lv) $ FunCall (genLoc storeVar) (dest : index ++ values)
	in genLoc $ Assign [lv'] [funC]

extractLVal :: LVal Var -> (Var, [TLExpr Var])
extractLVal lv = case lv of
    LVVar lvar    -> (unLoc lvar, [])
    LVStruct lv' fld -> 
        mapSnd (structAccess Bullet fld :) (extractLVal lv')
    LVCont ty lv' apat ->
        mapSnd (extractAPat ty apat :) (extractLVal lv')
    where
    extractAPat :: Type Var -> APat Var -> TLExpr Var
    extractAPat ty (VectP (CElem i)) = vectorAccess ty i
    extractAPat ty (VectP (CRange i j)) = vectorRange ty i j
    extractAPat ty (MatP (CElem i) (CElem j)) = matrixAccess ty i j
    extractAPat ty (MatP (CRange i j) (CRange k l)) = matrixRange ty i j k l
    extractAPat ty (MatP (CRange i j) (CElem k)) = matrixRowRange ty i j k
    extractAPat ty (MatP (CElem i) (CRange j k)) = matrixColRange ty i j k

    -- TODO: Are these type annotations correct?
    structAccess ty v = genLoc $ annTyE ty $ FunCall (genLoc $ sfield ty) [ genLoc $ annTyE (varType v) $ Var v ]
    vectorAccess ty v = genLoc $ annTyE ty $ FunCall (genLoc $ vind ty) [v]
    vectorRange ty v1 v2 = genLoc $ annTyE ty $ FunCall (genLoc $ vrange ty) [v1, v2]
    matrixAccess ty v1 v2 =      genLoc $ annTyE ty $ FunCall (genLoc $ mind ty) [v1, v2]
    matrixRange ty v1 v2 v3 v4 = genLoc $ annTyE ty $ FunCall (genLoc $ mrange ty)
                                                  [v1, v2, v3, v4]
    matrixColRange ty v1 v2 v3 = genLoc $ annTyE ty $ FunCall (genLoc $ mcolrange ty)
                                                  [v1, v2, v3]
    matrixRowRange ty v1 v2 v3 = genLoc $ annTyE ty $ FunCall (genLoc $ mrowrange ty)
                                                  [v1, v2, v3]
--
--        
--------------------------------------------------------------------------------
-- Removes all function and procedure calls due to global variables or
-- structured type accesses.

variableId :: LExpr Var -> Var
variableId (unLoc -> Var v) = v
variableId _                = error "<SSA.hs>.<variableId>: unexpected expr"

removeLoadStore :: CaoCFG -> CaoCFG
removeLoadStore cfg = cfg { blocks = blks }
    where 
    blks = Map.map (mapFst (concatMap (renameGVars (getWVars cfg) . aux))) (blocks cfg)
    aux :: LStmt Var -> [LStmt Var]
    aux ss@(unLoc -> FCallS fn _)
        | isStoreGlobal fn = [] -- Global variable store
        | otherwise        = [ss]
    aux ss@(unLoc -> Assign lv [unLoc -> TyE tyann (FunCall (unLoc -> fn) args)])
        | isStoreInit fn
            = let lvv = lvname $ head lv
                  ty  = varType lvv
              in [L (getLoc ss) $ VDecl $ ContD (genLoc lvv) (type2TyDecl ty) args]
        | isLoadGlobal fn = []
        | isStoreVar fn
            = let fstElem  = head args
                  lastElem = last args
                  lVal     = init $ tail args
              in [ L (getLoc ss) $ Assign [restoreLVal fstElem lVal] [lastElem]
                 , L (getLoc ss) $ Assign lv [fstElem]
                 ]
        | isLoadStruct fn
            = [ L (getLoc ss) $
                Assign lv [genLoc $ TyE tyann $ StructProj (head args)
                                               (variableId (unTypL (args!!1)))] ]
        | isLoadVar fn
            = [ L (getLoc ss) $
                Assign lv [ genLoc $ TyE tyann $ -- TODO: Verify TyE annotations
                                                 -- before: (annTy $ queryLVTy $ head lv)
                            Access (head args)
                                   (VectP (CElem (args!!1))) ] ]
        | isLoadVarRange fn
            = [ L (getLoc ss) $
                Assign lv [ genLoc $ TyE tyann $
                            Access (head args)
                                   (VectP (CRange (args!!1) (args!!2))) ] ]
        | isLoadMat fn
            = [ L (getLoc ss) $
                Assign lv [ genLoc $ TyE tyann $
                            Access (head args)
                                   (MatP (CElem (args!!1))
                                         (CElem (args!!2))) ] ]
        | isLoadMatRange fn
            = [ L (getLoc ss) $
                Assign lv [ genLoc $ TyE tyann $
                            Access (head args)
                                   (MatP (CRange (args!!1) (args!!2))
                                         (CRange (args!!3) (args!!4))) ] ]
        | isLoadMatRowR fn
            = [ L (getLoc ss) $
                Assign lv [ genLoc $ TyE tyann $
                            Access (head args)
                                   (MatP (CRange (args!!1) (args!!2))
                                         (CElem (args!!3))) ] ]
        | isLoadMatColR fn
            = [ L (getLoc ss) $
                Assign lv [ genLoc $  TyE tyann $
                            Access (head args)
                                   (MatP (CElem (args!!1))
                                         (CRange (args!!2) (args!!3))) ] ]
    aux ss
          = [ ss ]

    renameGVars :: [Var] -> BasicBlock -> BasicBlock
    renameGVars wvars = map (rnGVars wvars)
    
    rnGVars :: [Var] -> LStmt Var -> LStmt Var
    rnGVars wvs (L l s) = L l $ fmap (rnGVars_ wvs) s
        
    rnGVars_ :: [Var] -> Var -> Var
    rnGVars_ wvs v
--          | Just v' <- find ((== varName v) . varName) wvs, Global <- varScope v
--          TODO: Check this function!!! Hack: rename variables except phi functions.
        | Just v' <- find ((== varName v) . varName) wvs, Global <- varScope v, not (isPhiFun v)
            = v'
        | otherwise
            = v

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

restoreLVal :: TLExpr Var -> [TLExpr Var] -> LVal Var
restoreLVal lvar [] = LVVar (L (getLoc lvar) $ variableId $ unTypL lvar)
restoreLVal lvar (x:xs) = aux x
    where 
    lv = restoreLVal lvar xs
    aux :: TLExpr Var -> LVal Var
    aux (unLoc -> unTyp -> FunCall (unLoc -> n) args)
        | isLValSField  n = LVStruct lv $ variableId $ unTypL $ head args
        | isLValVInd    n = LVCont (varType n) lv $ VectP $ CElem $ head args
        | isLValVRng    n = LVCont (varType n) lv $ VectP $ CRange (head args)
                                                                   (args!!1)
        | isLValMInd    n = LVCont (varType n) lv $ MatP (CElem (head args))
                                                         (CElem (args!!1))
        | isLValMRng    n = LVCont (varType n) lv $
                                MatP (CRange (head args) (args!!1))
                                     (CRange (args!!2) (args!!3))
        | isLValMColRng n = LVCont (varType n) lv $
                                MatP (CElem (head args))
                                     (CRange (args!!1) (args!!2))
        | isLValMRowRng n = LVCont (varType n) lv $
                                MatP (CRange (head args) (args!!1))
                                     (CElem (args!!2))
    aux _ 
          = error "<SSA>.<restoreLVal>: unexpected case"


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


toSSA :: CaoMonad m => CaoCFG -> m CaoCFG
-- Ignore empty blocks (type and variable definitions)
toSSA cfg | Map.null (blocks cfg) = return cfg
          | otherwise             = renameVars dt blocks' vars
    where
    g               = graphFromEdges_ $ blocks cfg
    dt              = genDomTree g
    cfg'            = introLoadStore cfg
    (blocks', vars) = insertPhiFuncs g cfg'

fromSSA :: CaoCFG -> CaoCFG
fromSSA cfg
    | Map.null (blocks cfg) = cfg
    | otherwise             = BT.fromSSA $ removeSsaDecl $ removeLoadStore cfg