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

{-
Module      :  $Header$
Description :  Translating back from 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.SsaBack
    ( fromSSA
    , introduceDefs
    , rmVars
    ) where

import Data.List ( foldl', partition, nubBy, insertBy, minimumBy )
import Data.Map ( Map )
import qualified Data.Map as Map
import Data.Set ( Set )
import qualified Data.Set as Set

import Language.CAO.Analysis.CFG
import Language.CAO.Syntax
import Language.CAO.Syntax.Utils
import Language.CAO.Type.Utils

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

fromSSA :: CaoCFG -> CaoCFG
fromSSA = introduceDefs . rmVars . coalescePhiFuncs
--      . eliminateInterference
--
--eliminateInterference :: CaoCFG -> CaoCFG
--eliminateInterference cfg = cfg
--  where liveRanges = undefined
--
--        interVars :: [Set String]
--        interVars = undefined
--
--        insertCopies :: Set String -> CaoCFG -> CaoCFG
--        insertCopies = undefined

-- type Location = (NodeId, Int)
-- type RLoc = Map String Location
-- 
-- type LiveIn = Map NodeId (Set String)
-- type LiveOut = Map NodeId (Set String)
-- 
-- type PhiCong = Map String (Set String)


coalescePhiFuncs :: CaoCFG -> CaoCFG
coalescePhiFuncs cfg = cfg { blocks = coalesceVars phiCong stmtMap }
    where 
    stmtMap :: LocalGraph
    stmtMap = blocks cfg

    phiCong :: Map Var (Var, Set Var)
    phiCong = Map.fold (\(stmts, _) m ->
                           foldl' getPhis m stmts)
                       Map.empty stmtMap

getPhis :: Map Var (Var, Set Var)
        -> LStmt Var
        -> Map Var (Var, Set Var)
getPhis m (unLoc -> Assign [LVVar (unLoc -> str)]
                           [unLoc -> unTyp -> FunCall (unLoc -> fId) args])
    | isPhiFun fId
        = fixPhiC m $ str : (Set.toList $ fvs args)
getPhis m _ = m

fixPhiC :: Map Var (Var, Set Var)
        -> [Var]
        -> Map Var (Var, Set Var)
fixPhiC m vs = Set.fold (\v -> Map.insert v congr) m congs
    where 
    congs = Set.unions $ map fGetPhiCong vs
    congr = (Set.findMin congs, congs)
    fGetPhiCong v = maybe (Set.singleton v) snd (Map.lookup v m)

coalesceVars :: Map Var (Var, Set Var)
             -> LocalGraph
             -> LocalGraph
coalesceVars cong = Map.map (coalesceVarsStmts rnEnv)
    where 
    rnEnv :: SEnv Var Var
    rnEnv = Map.foldWithKey renameF emptyRN cong

    -- All variables in a phi congruence class are renamed to
    -- one "representative" element (this case, the head of the list)
    renameF :: Var -> (Var, Set Var) -> SEnv Var Var -> SEnv Var Var
    renameF v1 (v2, _) b = b +> v1 ~> v2

-- XXX: can the order of composition be changed, ie., the filter can be
-- performed before the map?
coalesceVarsStmts :: (SEnv Var Var)
                  -> (BasicBlock, Connections)
                  -> (BasicBlock, Connections)
coalesceVarsStmts cong = mapFst (filter (not . isPhi . unLoc) . map (<|> cong))
    where 
    isPhi :: Stmt Var -> Bool
    isPhi (Assign [LVVar _] [unLoc -> unTyp -> FunCall fId _])
        = isPhiFun (unLoc fId)
    isPhi _                  = False


---------------------------------------------------------
-- TODO: REFACTOR vars <---> except (almost the same def)
---------------------------------------------------------

introduceDefs :: CaoCFG -> CaoCFG
introduceDefs cfg = cfg { blocks = addDecls (map mkDecl $ sortDeps neededDefs) blk }
--introduceDefs cfg = addDecls (map mkDecl $ sortDeps neededDefs) cfg
    where 
    blk :: LocalGraph
    blk = blocks cfg

    neededDefs :: [Var]
    neededDefs = filter isLocal $ Set.toList $
            vars `Set.difference` except

    vars, except, alreadyDef, args :: Set Var
    vars       = Map.fold foldVars Set.empty blk
    except     = alreadyDef `Set.union` args
    alreadyDef = Map.fold foldDecls Set.empty blk
    args       = bvs $ definition cfg

    foldDecls :: (BasicBlock, Connections) -> Set Var -> Set Var
    foldDecls (stmts, _) s0 = s0 `Set.union` bvs stmts

    foldVars  :: (BasicBlock, Connections) -> Set Var -> Set Var
    foldVars  (stmts, _) s0 = s0 `Set.union` fvs stmts

--addDecls :: [LStmt Var] -> CaoCFG -> CaoCFG
addDecls :: [LStmt Var] -> LocalGraph -> LocalGraph
addDecls lst = Map.alter fAddDecl (entryNode + 1)
    where
    fAddDecl :: Maybe (BasicBlock, Connections) 
             -> Maybe (BasicBlock, Connections)
    fAddDecl = fmap (mapFst (lst ++))
  
mkDecl :: Var -> LStmt Var
mkDecl v = genLoc $ VDecl $ VarD (genLoc v) (type2TyDecl $ varType v) Nothing

sortDeps :: [Var] -> [Var]
sortDeps = sortDeps' Set.empty
  where sortDeps' _  [] = []
        sortDeps' ds vs = vs1 ++ sortDeps' (ds `Set.union` Set.fromList vs1) vs2
          where (vs1, vs2) = partition noDeps vs
                noDeps v   = Set.filter isLocal (fvs $ varType v)
                                `Set.isSubsetOf` ds

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

rmVars :: CaoCFG -> CaoCFG
rmVars cfg = CaoCFG { definition = d0, blocks = b0 }
    where 
    rnMap :: SEnv Var Var
    rnMap = mkRenameMap $ varRange cfg

    b0 :: LocalGraph
    b0 = filterDecls $ Map.map rmAndFilter $ blocks cfg

    rmAndFilter :: (BasicBlock, Connections) -> (BasicBlock, Connections)
    rmAndFilter = mapFst (\ stmts -> filter filterAssigns $ stmts <|> rnMap)

    d0 :: LDef Var
    d0 = fmap (fmap (<|> rnMap)) (definition cfg)

    filterAssigns :: LStmt Var -> Bool
    filterAssigns (unLoc -> Assign [LVVar v] [unLoc -> unTyp -> Var v'])
        = unLoc v /= v'
    filterAssigns _
        = True

filterDecls :: LocalGraph -> LocalGraph
filterDecls = Map.map (mapFst nubDecls)
    where 
    nubDecls :: [LStmt Var] -> [LStmt Var]
    nubDecls = nubBy eqDecls

    eqDecls :: LStmt Var -> LStmt Var -> Bool
    eqDecls (L _ (VDecl v0)) (L _ (VDecl v1))
          = eqVarDecls v0 v1
    eqDecls _ _
          = False

        -- TODO: INCOMPLETE DEFINITION (MAY CAUSE BUGS?)
    eqVarDecls :: VarDecl Var -> VarDecl Var -> Bool
    eqVarDecls (VarD   v0 _ _) (VarD v1 _ _ ) = v0 == v1
    eqVarDecls (MultiD v0 _  ) (MultiD v1 _ ) = v0 == v1
    eqVarDecls (ContD  v0 _ _) (ContD v1 _ _) = v0 == v1
    eqVarDecls _               _              = False

mkRenameMap :: VarRange -> SEnv Var Var
mkRenameMap rng = foldl' (+>) emptyRN $ map mkM gRanges
  
    where 
    rnglst :: [(Var, Range)]
    rnglst = Map.toList rng

    -- vars grouped same type
    gTypes :: [[(Var, Range)]]
    gTypes = groupType rnglst

    gRanges :: [[(Var, Range)]]
    gRanges = concatMap (accumRanges [] []) gTypes

    mkM :: [(Var, Range)] -> SEnv Var Var
    mkM []  = emptyRN
    mkM [_] = emptyRN
    mkM xs  = let (x, _) = minimumBy cmpRng xs
              in foldl' (\a (b, _) -> a +> b ~> x) emptyRN (init xs)

    accumRanges :: [(Var, Range)]
                    -> [(Var, Range)]
                    -> [(Var, Range)]
                    -> [[(Var, Range)]]
    accumRanges [] acc []   = [acc]
    accumRanges orig acc [] = acc:accumRanges [] [] orig
    accumRanges orig acc (x:xs) = case mutuallyDisj x acc of
        Just lst -> accumRanges orig lst xs
        Nothing  -> accumRanges (x:orig) acc xs

    mutuallyDisj :: (Var, Range)
                     -> [(Var, Range)]
                     -> Maybe [(Var, Range)]
    mutuallyDisj l@(v, r0) rg
          | not (isContainer (varType v)) && all (disjoint r0 . snd) rg
              = Just (l:rg)
          | lst <- insertBy cmpRng l rg, chainsSafely lst
              = Just lst
          | otherwise
              = Nothing

    cmpRng :: (Var, Range) -> (Var, Range) -> Ordering
    cmpRng (_, FromTo l00 l01 _ _) (_,FromTo l10 l11 _ _)
          | l10 `gtLoc` l01 = LT
          | l00 `gtLoc` l11 = GT
          | otherwise       = compare l00 l10

    chainsSafely :: [(Var, Range)] -> Bool
    chainsSafely []
          = True
    chainsSafely [_]
          = True
    chainsSafely ((_, FromTo _ l0 _ _):rest@((_, FromTo l1 _ ab1 _):_))
          = l1 `gtLoc` l0 && ab1 && chainsSafely rest

    groupType :: [(Var, Range)]
                  -> [[(Var, Range)]]
    groupType [] = []
    groupType (x:rest)
          = let (st, r) = partition (sameType x) rest
            in (x:st):groupType r

    sameType :: (Var, Range)
                 -> (Var, Range)
                 -> Bool
    sameType (t0, _) (t1, _)
          = varType t0 == varType t1

    disjoint :: Range -> Range -> Bool
    disjoint (FromTo l00 l01 _ _) (FromTo l10 l11 _ _) =
        l11 `gtLoc` l10 && l01 `gtLoc` l00 && 
            (l10 `gtLoc` l01 || l00 `gtLoc` l11)

    gtLoc :: Location -> Location -> Bool
    gtLoc (n0, loc0) (n1, loc1)
        -- Special case when 0
        | n1 == 0                  = False
        | n0 == 0                  = True
        -- Lexicographic order
        | n0 > n1                  = True
        | n0 < n1                  = False
        -- When n1 == n2
        | otherwise                = loc0 >= loc1

type Location = (NodeId, Int)
data Range    = FromTo { _fromL :: Location
                       , _toL   :: Location
                       , _safeA :: Bool    -- safeA is True when the variable is
                                           -- initialized completely in its
                                           -- first assignment
                       , _safeL :: Bool    -- if safeL = True, it is safe to
                                           -- consider >= instead of > to check
                                           -- disjoint live ranges,
                       } deriving Show
type VarRange = Map Var Range

-- NOTE: USED ONLY FOR TESTING PURPOSES!
--showVR m = Map.foldWithKey (\k a acc -> showPpr k ++ "\\\\\\" ++ show a ++ "\n" ++ acc) "" m

varRange :: CaoCFG -> VarRange
varRange cfg = vRange
    where 
    vRange :: VarRange
    vRange = Map.filterWithKey (\k _ -> isLocal k && not (k `elem` seqVars))
                                gvRange
        
    gvRange :: VarRange
    gvRange = traverseCFG [entryNode + 1] [] argsRange (blocks cfg)

    seqVars :: [Var]
    seqVars = getSeqVars cfg

    argsRange :: VarRange
    argsRange = Set.fold fArgs Map.empty $ bvs $ definition cfg

    fArgs :: Var -> VarRange -> VarRange
    fArgs v m
        | nsVar v   = Map.insert v (FromTo (entryNode, 0)
                                           (exitNode, 0)
                                           False
                                           False) m
        | otherwise = m

getSeqVars :: CaoCFG -> [Var]
getSeqVars = concatMap doGetSV . concatMap fst . Map.elems . blocks
    where 
    doGetSV (unLoc -> Seq i _) = [seqVar i]
    doGetSV _                  = []

-- NOTE: if a variable is first assigned in node 3, loc 5, but
-- there is a loop back to node 2, its range should be fixed to be from node 2,
-- loc 0:
--   while ..
--      ... x1
--      ...
--      x1 := ...
-- This should not be a problem, as the first x1 will never be renamed. But
-- we should be careful about this.
traverseCFG :: [NodeId]
            -> [NodeId]
            -> VarRange
            -> LocalGraph 
            -> VarRange
traverseCFG []     _    m _
  = m
traverseCFG (n:ns) seen m cfg
    | n `elem` seen = traverseCFG ns seen m cfg
    | otherwise     = traverseCFG (ns ++ next) (n:seen) m' cfg
    where 
    blk  :: BasicBlock
    next :: Connections
    (blk, next) = cfg Map.! n

    m' :: VarRange
    m' = foldl' updateRanges m nStmts
        
    nStmts :: [(Location, LStmt Var)]
    nStmts = zip [ (n, i) | i <- [1..] ] blk

updateRanges :: VarRange -> (Location, LStmt Var) -> VarRange
updateRanges rng (lloc,ss@(unLoc -> Assign lvs _))
    | all isSimpleLVal lvs
        = fixAssignRng True  lvns rvns
    | otherwise
        = fixAssignRng False lvns rvns
    where 
    fixAssignRng :: Bool -> [Var] -> [Var] -> VarRange
    fixAssignRng safe lVars rVars
          = let rng'  = foldl' (fixLRanges safe lloc) rng lVars
            in foldl' (fixRanges safe lloc) rng' rVars
    lvns = lvnames ss
    rvns = lvns ++ rvnames ss
updateRanges rng (lloc, ss)
    = foldl' (fixRanges False lloc) rng (rvnames ss)

fixLRanges :: Bool -> Location -> VarRange -> Var -> VarRange
fixLRanges b l = flip (Map.alter (updateDefRange b l))

fixRanges :: Bool -> Location -> VarRange -> Var -> VarRange
fixRanges b l  = flip (Map.alter (updateLastRange b l))

updateDefRange :: Bool -> Location -> Maybe Range -> Maybe Range
updateDefRange b loc Nothing = Just (FromTo loc (exitNode, 0) b False)
updateDefRange ab l@(n, loc) mr@(Just (FromTo (n0,l0) loc1 _ b))
    | n0 < n || (n0 == n && l0 < loc) = mr
    | otherwise                       = Just $ FromTo l loc1 ab b

updateLastRange :: Bool -> Location -> Maybe Range -> Maybe Range
-- If it was not previously used as a lvalue, it should be considered
-- alive along the whole CFG
updateLastRange _ _   Nothing = Just (FromTo (entryNode, 0)
                                             (exitNode , 0)
                                             False
                                             False)
updateLastRange b l@(n,loc) mr@(Just (FromTo loc1 (n0,l0) ab _))
    | n0 > n || (n0 == n && l0 > loc) = mr
    | otherwise                       = Just $ FromTo loc1 l ab b

-- TODO: REFACTOR IN Language.CAO.Syntax.Utils
lvnames :: LStmt Var -> [Var]
lvnames (unLoc -> Assign lvs _) = map lvname lvs
lvnames _                       = []

rvnames :: LStmt Var -> [Var]
rvnames (unLoc -> ss@(Assign lvs _))
    = Set.toList rvs
    where 
    vs     = fvs ss
    vlvs   = Set.fromList $ map lvname lvs
    rvs    = vs Set.\\ vlvs

rvnames stmt
    = Set.toList $ fvs stmt