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

{- |
    Module      :  $Header$
    Description :  CAO control flow graph.
    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 control flow graph abstractions.
-}

module Language.CAO.Analysis.CFG 
    ( NodeId
    , BasicBlock
    , Connections
    , LocalGraph
    , CaoCFG(..)
    , buildCFG
    , toAST
    , getDefFromBlocks
    , swaps
    , showCFG
    , entryNode
    , exitNode
    , removeSsaDecl
    , graphFromEdges_
    ) where

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

import Language.CAO.Analysis.Dominance

import Language.CAO.Common.Outputable
import Language.CAO.Common.SrcLoc
import Language.CAO.Common.Utils
import Language.CAO.Common.Var

import Language.CAO.Syntax
import Language.CAO.Syntax.Utils (isSimpleVDecl)

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

type NodeId      = Int
type BasicBlock  = [LStmt Var]
type Connections = [NodeId]
type LocalGraph  = Map NodeId (BasicBlock, Connections)

data CaoCFG = CaoCFG {
    definition :: LDef Var,
    blocks     :: LocalGraph
 } 

data CFGState = CFGState {
    currentId   :: !NodeId,
    currentNode :: BasicBlock,
    graph       :: LocalGraph
 } 

emptyState :: CFGState
emptyState = CFGState (entryNode + 1) [] $ 
    Map.fromList [ (entryNode, ([], [])), (exitNode, ([], [])) ]

entryNode :: NodeId
entryNode = 1

exitNode :: NodeId
exitNode = 0

currentNodeId :: State CFGState NodeId
currentNodeId = gets currentId

addToCurrentNode :: LStmt Var -> State CFGState ()
addToCurrentNode stmt = 
    modify $ \ s -> s { currentNode = stmt : currentNode s }

endCurrentNode :: State CFGState ()
endCurrentNode = modify $ \ s -> let
        cnode = currentNode s
    in if null cnode
        then s 
        else s { graph = Map.insert (currentId s) (reverse cnode, []) (graph s),
                 currentNode = [],
                 currentId = succ (currentId s) }

addEmptyNode :: State CFGState NodeId
addEmptyNode = get >>= \ s -> do
    let i = currentId s
        cnode = reverse $ currentNode s
    put $ s { graph = Map.insert i (cnode, []) (graph s),
                 currentNode = [],
                 currentId = succ i
                }
    return i

-- Precondition: endCurrentNode should have been called first
addSingleNode :: LStmt Var -> State CFGState NodeId
addSingleNode stmt = do
    addToCurrentNode stmt
    cid <- currentNodeId
    endCurrentNode
    return cid

addEdges :: [NodeId] -> [NodeId] -> State CFGState ()
addEdges origin target = mapM_ (`addEdge` target) origin

-- This needs to add target nodes to the end of the list,
-- in order to ensure the correct order when translating
-- back to an AST
addEdge :: NodeId -> [NodeId] -> State CFGState ()
addEdge origin target = modify $ \ s -> 
    s { graph = Map.adjust (mapSnd (++ target)) origin (graph s) }

--------------------------------------------------------------------------------
-- Generation of CFG

--------------------------------------------------------------------------------
------------------------------------------ Prog --------------------------------

buildCFG :: Prog Var -> [CaoCFG]
buildCFG (Prog defs _) = map definitionCFG defs

--------------------------------------------------------------------------------
--------------------------------------- Definition -----------------------------

definitionCFG :: LDef Var -> CaoCFG
definitionCFG (L l (FunDef (Fun n args rt body)))
    | [] <- filter (not . isSimpleVDecl . unLoc) body  =
        let st = flip execState emptyState $ do
                 nd <- addEmptyNode
                 addEdges [entryNode] [nd]
                 addEdges [nd] [exitNode]
        in CaoCFG fd (graph st)
    | otherwise =
        let st = flip execState emptyState $ do
                 (entry, exit) <- toGraph body
                 addEdges [entryNode] entry
                 addEdges exit [exitNode]
        in CaoCFG fd (graph st)
    where 
    fd = L l $ FunDef $ Fun n args rt []
definitionCFG def
  = CaoCFG def Map.empty


--------------------------------------------------------------------------------
------------------------------------------ Func --------------------------------

-- The function returns the id of the entry node of a sub-graph and the list
-- of id's of exit blocks of a sub-graph
toGraph :: [LStmt Var] -> State CFGState ([NodeId], [NodeId])
toGraph [] = do
    curNode <- currentNodeId
    endCurrentNode
    return ([curNode], [curNode])
toGraph (L l (VDecl (ContD vn _ exs)) : xs) = do
    let sv = L (getLoc vn) $ mkStoreInit $ varName $ unLoc vn
    declEntry l vn sv exs
    toGraph xs
-- Ignoring variable declaration
toGraph (L l (VDecl (VarD vn _ _)) : xs) = do
    declEntrySSA l vn
    toGraph xs
toGraph (L l (VDecl (MultiD vns _)) : xs) = do
    mapM_ (declEntrySSA l) vns
    toGraph xs
toGraph (L l (CDecl (ConstD cn _ _)) : xs) = do
    declEntrySSA l cn
    toGraph xs
toGraph (L l (CDecl (MultiConstD cns _ _)) : xs) = do
    mapM_ (declEntrySSA l) cns
    toGraph xs
toGraph (s@(L _ (Assign _ _)) : xs)           =
    addToCurrentNode s >> toGraph xs
toGraph (s@(L _ (FCallS _ _)) : xs) =
    addToCurrentNode s >> toGraph xs
toGraph (s@(L _ (Ret _)) : _)          = do
    curNode <- currentNodeId
    addToCurrentNode s 
    endCurrentNode
    addEdges [curNode] [exitNode]
    return ([curNode], []) -- There is no exit point since the node is final
toGraph (L l (Ite i t Nothing) : xs) = do
    curNode <- currentNodeId
    addToCurrentNode (L l $ Ite i [] Nothing)
    endCurrentNode
    (ifEntryNode,   ifExitNodes)   <- toGraph t
    (nextEntryNode, nextExitNodes) <- toGraphRest xs
    addEdges [curNode] (ifEntryNode ++ nextEntryNode)
    addEdges ifExitNodes nextEntryNode
    return ([curNode], if null xs then curNode : ifExitNodes else nextExitNodes)
toGraph (L l (Ite i t (Just e)) : xs) = do
    curNode <- currentNodeId
    addToCurrentNode (L l $ Ite i [] (Just []))
    endCurrentNode
    (ifEntryNode,   ifExitNodes)   <- toGraph t
    (elseEntryNode, elseExitNodes) <- toGraph e
    (nextEntryNode, nextExitNodes) <- toGraphRest xs
    addEdges [curNode] (ifEntryNode ++ elseEntryNode)
    addEdges (ifExitNodes ++ elseExitNodes) nextEntryNode
    return ([curNode], 
            if null xs then ifExitNodes ++ elseExitNodes else nextExitNodes)
toGraph (L l (While cond wstmts) : xs) = do
    curNode <- currentNodeId
    endCurrentNode
    whileNode <- addSingleNode (L l $ While cond [])
    (bodyEntryNode, bodyExitNodes) <- toGraph wstmts
    (nextEntryNode, nextExitNodes) <- toGraphRest xs
    when (curNode /= whileNode) $ addEdges [curNode] [whileNode]
    addEdges [whileNode] (bodyEntryNode ++ nextEntryNode)
    addEdges bodyExitNodes [whileNode]
    return ([curNode], if null xs then [whileNode] else nextExitNodes)
toGraph (L l (Seq iter stmts) : xs) = do
    curNode <- currentNodeId
    endCurrentNode
    seqNode <- addSingleNode (L l $ Seq iter [])
    (bodyEntryNode, bodyExitNodes) <- toGraph stmts
    (nextEntryNode, nextExitNodes) <- toGraphRest xs
    when (curNode /= seqNode) $ addEdges [curNode] [seqNode]
    addEdges [seqNode] (bodyEntryNode ++ nextEntryNode)
    addEdges bodyExitNodes [seqNode]
    return ([curNode], if null xs then [seqNode] else nextExitNodes)
toGraph (s@(L _ (Nop _)) : xs) = 
    addToCurrentNode s >> toGraph xs

declEntrySSA :: SrcLoc -> Located Var -> State CFGState ()
declEntrySSA l vn = declEntry l vn (genLoc ssaDecl) []

declEntry :: SrcLoc -> Located Var -> Located Var -> [TLExpr Var] -> State CFGState ()
declEntry l v fc exs = addToCurrentNode $ 
    L l $ Assign [LVVar v] [ L l $ annTyE (varType $ unLoc v) $ FunCall fc exs]

toGraphRest :: [LStmt Var] -> State CFGState ([NodeId], [NodeId])
toGraphRest [] = return ([], [])
toGraphRest xs = toGraph xs
--------------------------------------------------------------------------------
-- Back from CFG

toAST :: [CaoCFG] -> Prog Var
toAST cfg = Prog (map getDefFromBlocks cfg) Nothing

getDefFromBlocks :: CaoCFG -> LDef Var
getDefFromBlocks (CaoCFG cdef cblocks) = case unLoc cdef of
    FunDef (Fun n args rt []) ->
        L (getLoc cdef) $ FunDef $ Fun n args rt $ bodyFromBlocks cblocks
    _ -> cdef

bodyFromBlocks :: LocalGraph -> [LStmt Var]
bodyFromBlocks blks = let
        g = graphFromEdges_ blks
        dt = invertMap $ genDomTree g
    in bodyFromBlocks' dt blks $ head $ snd $ blks Map.! entryNode

graphFromEdges_ :: LocalGraph -> Graph
graphFromEdges_ blks = let
        (g,_,_) = graphFromEdges $ map (\ (k, (_, c)) -> (k, k, c) ) $ Map.assocs blks
    in g

bodyFromBlocks' :: Map Vertex [Vertex] -> LocalGraph -> NodeId -> [LStmt Var]
bodyFromBlocks' domTree blks nid = 
    if nid == exitNode then [] else let
        nextNodes = domTree Map.! nid
        (bn, cn) =  blks    Map.! nid
        (stmts, lastStmt) = initLast bn
    in if null bn then [] else
        case unLoc lastStmt of
            Ite i _ Nothing  -> stmts ++ ( L (getLoc lastStmt) 
                    (Ite i (fetchNextBlock (cn !! 0)) Nothing) :
                fetchNext (nextNodes \\ [cn !! 0]))
            Ite i _ (Just _) -> stmts ++ ( L (getLoc lastStmt) 
                    (Ite i (fetchNextBlock (cn !! 0))
                           (Just $ fetchNextBlock (cn !! 1))) : 
                fetchNext (nextNodes \\ [cn !! 0, cn !! 1]))
            While c _        -> stmts ++ ( L (getLoc lastStmt)
                    (While c (fetchNextBlock (cn !! 0))) :
                fetchNext (nextNodes \\ [cn !! 0]))
            Seq i _          -> stmts ++ ( L (getLoc lastStmt)
                    (Seq i (fetchNextBlock (cn !! 0) )) :
                fetchNext (nextNodes \\ [cn !! 0]))
            _                -> bn    ++ fetchNext nextNodes
    where
    fetchNextBlock = bodyFromBlocks' domTree blks
    fetchNext = concatMap (bodyFromBlocks' domTree blks)

--- auxiliary ---

showCFG :: [CaoCFG] -> String
showCFG cfg = "digraph cao_cfg {\n" ++ unlines (map aux cfg) ++ invisedgs ++ "\n}"
    where

    invisedgs
      | null invisedgs' = ""
      | otherwise       = "edge [style = invis]\n" ++ unlines invisedgs'

    invisedgs' = graphs cfg

    graphs :: [CaoCFG] -> [String]
    graphs []  = []
    graphs [_] = []
    graphs ((CaoCFG (unLoc->FunDef f1) _):c@(CaoCFG (unLoc->FunDef f2) _):rest)
      = let edg = "\"0" ++ showPprIds (funId f1)
                  ++ "\" -> \"1" ++ showPprIds (funId f2) ++ "\"\n"
        in  edg:(graphs (c:rest))
    graphs (d@(CaoCFG (unLoc->FunDef _) _):_:rest)
      = graphs (d:rest)
    graphs (_:rest)
      = graphs rest

    aux :: CaoCFG -> String
    aux (CaoCFG def bk) = case unLoc def of
        FunDef f -> let fundef = showPprIds (funId f)
                    in unlines $ 
            ["subgraph " ++ fundef ++ " {\n"] ++ nodes fundef bk ++ cfgEdges fundef bk ++ 
            ["}"]
        _ -> ""

    nodes str = map (\ (k, (s, _)) -> let 
            sst = if k == entryNode then "Entry\\l" ++ str ++ "\\n" else if k == exitNode then "Exit\\n" else ""
        in "node [label=\"" ++ sst ++ showStmts s ++ "\"]\n\"" ++ show k  ++ str ++ "\" [shape=box];")
        . Map.assocs
    cfgEdges str = concatMap (\ (k, (_, nl)) -> 
        map (\b -> '"':show k ++ str ++ "\" -> \"" ++ show b ++ str ++ "\"") nl) . Map.assocs

    showStmts :: PP a => [a] -> String
    showStmts = concatMap ((++ "\\l") . filter (/= '\n') . showPprIds)

removeSsaDecl :: CaoCFG -> CaoCFG
removeSsaDecl cfg = cfg { blocks = Map.map filterSsaDecls (blocks cfg) }
  where filterSsaDecls :: (BasicBlock, Connections) -> (BasicBlock, Connections)
        filterSsaDecls (ss, n) = (filter (not . isSsaDeclStmt . unLoc) ss, n)

        isSsaDeclStmt :: Stmt Var -> Bool
        isSsaDeclStmt (Assign [LVVar _] [unLoc -> unTyp -> FunCall fn []]) = isSsaDecl $ unLoc fn
        isSsaDeclStmt _                                           = False