{- 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 . -}
{-# 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
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 ".:\
\ 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 ".: 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 ".: 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 ".: 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 ".: 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