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

{-
Module      :  $Header$
Description :  Insertion of phi function in SSA 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

-}
module Language.CAO.Analysis.PhiInsert where

import Data.Graph ( Graph, Vertex )
import Data.List
import Data.Map ( Map )
import qualified Data.Map as Map
import Data.Maybe ( fromMaybe )
import Data.Set ( Set )
import qualified Data.Set as Set

import Language.CAO.Common.Var
import Language.CAO.Common.SrcLoc

import Language.CAO.Syntax
import Language.CAO.Syntax.Utils hiding ( getVars )

import Language.CAO.Analysis.Dominance
import Language.CAO.Analysis.CFG


--To decide what blocks require a phi
--function to join a definition to a variable
--v in block b:
--  1. Compute D1 = DF(b).
--  Place Phi functions at the head of all
--  members of D1.
--  
--  2. Compute D2 = DF(D1).
--  Place Phi functions at the head of all
--  members of D2-D1.
--  
--  3. Compute D3 = DF(D2).
--  Place Phi functions at the head of all
--  members of D3-D2-D1.
--  
--  4. Repeat until no additional Phi
--  functions can be added.

insertPhiFuncs :: Graph -> CaoCFG -> (CaoCFG, [Var])
insertPhiFuncs g cfg = (phiIns phiFuns cfg, allVars)
    where
    phiFuns :: Set (Vertex, Var, Int)
    phiFuns = phiLocArity g locOf
    allVars :: [Var]
    allVars = Map.keys locOf
    locOf   :: Map Var (Set Vertex)
    locOf   = Map.foldWithKey getLocs Map.empty (blocks cfg)

getLocs :: Vertex -> ([LStmt Var], [NodeId]) -> Map Var (Set Vertex)
        -> Map Var (Set Vertex)
getLocs nI (stmts, _) lo = foldl' (addVerts nI) lo $ getVars stmts

addVerts :: Vertex -> Map Var (Set Vertex) -> Var -> Map Var (Set Vertex)
addVerts nI lo v = Map.alter (addVertex nI) v lo

addVertex :: Vertex -> Maybe (Set Vertex) -> Maybe (Set Vertex)
addVertex v Nothing  = Just $ Set.singleton v
addVertex v (Just s) = Just $ Set.insert v s
        

phiLocArity :: Graph -> Map Var (Set Vertex) -> Set (Vertex, Var , Int)
phiLocArity g = Map.foldWithKey foldDf Set.empty
    where
    foldDf :: Var -> Set Vertex -> Set (Vertex, Var, Int)
           -> Set (Vertex, Var, Int)
    foldDf s v acc
        | Set.size v > 1 =
            Set.map (phiVarArity s) (followDf Set.empty v) `Set.union` acc
        | otherwise      = acc

    phiVarArity :: Var -> Vertex -> (Vertex, Var, Int)
    phiVarArity s v = (v, s, length $ predecessors g v)

    followDf :: Set Vertex -> Set Vertex -> Set Vertex
    followDf ini d1
        | d2 <- Set.fold getDF ini d1, d1 /= d2 = followDf d2 d2
        | otherwise                               = d1

    getDF :: Vertex -> Set Vertex -> Set Vertex
    getDF v s0 = Set.union s0 $ fromMaybe Set.empty $ Map.lookup v df

    df :: Map Vertex (Set Vertex)
    df = domFront g

phiIns :: Set (Vertex, Var, Int) -> CaoCFG -> CaoCFG
phiIns s cfg | Set.size s == 0 = cfg
             | otherwise       = phiIns s' cfg'
    where
    (phiAt@(_, n, _), s')  = Set.deleteFindMax s
    phiFunN                = mkPhiFunVar (varName n)
    cfg'                   = phiIns_ phiFunN phiAt cfg

phiIns_ :: Var -> (Vertex, Var, Int) -> CaoCFG -> CaoCFG
phiIns_ phiFunN (nodeI, vname, arity) cfg
    | nodeI == exitNode = cfg
    | otherwise         = addStmtAt nodeI phiFunAssign cfg
    where
    phiFunAssign :: LStmt Var
    phiFunAssign = genLoc $ Assign [lval] [phiFun]
    lval :: LVal Var
    lval = LVVar (genLoc vname)
    phiFun :: TLExpr Var
    -- XXX: Is this the correct annotation type?
    phiFun = genLoc $ annTyE (varType vname) $ FunCall (genLoc phiFunN) args
    args :: [TLExpr Var]
    args = map (genLoc . annTyE (varType vname) . Var) $ replicate arity vname

addStmtAt :: Vertex -> LStmt Var -> CaoCFG -> CaoCFG
addStmtAt nodeI stmt cfg = cfg { blocks = Map.insert nodeI (stmt:stmts, c) blk }
    where
    blk       = blocks cfg
    (stmts,c) = blk Map.! nodeI

getVars :: [LStmt Var] -> [Var]
getVars = concatMap (variableName . unLoc)
    where
    variableName :: Stmt Var -> [Var]
    variableName (VDecl v)          = Set.toList $ bvs v
    variableName (Assign lvalues _) = Set.toList $ fvs lvalues
    variableName _                  = []