{-# LANGUAGE DeriveGeneric #-}


module Data.CFG where


import Data.X86.Instruction

import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Set as S
import Data.Maybe


import qualified Data.Serialize as Cereal hiding (get,put)
import Control.DeepSeq
import GHC.Generics


-- | A control flow graph with blocks and edges.
-- A blockID (represented as an @Int@) is a unique identifier of a basic block.
-- We store basic blocks twice: once as addresses, and once as instructions.
data CFG = CFG {
  CFG -> IntMap [Int]
cfg_blocks :: IM.IntMap [Int],            -- ^ A mapping of blockIDs to instruction addresses
  CFG -> IntMap IntSet
cfg_edges  :: IM.IntMap (IS.IntSet),      -- ^ A mapping of blockIDs to sets of blocKIDs
  CFG -> IntMap Int
cfg_addr_to_blockID :: IM.IntMap Int,     -- ^ A mapping of instruction addresses to blockIDs
  CFG -> Int
cfg_fresh :: Int,                         -- ^ A fresh blockID
  CFG -> IntMap [Instruction]
cfg_instrs :: IM.IntMap [Instruction]     -- ^ A mapping of blockIDs to instructions
 }
 deriving (Int -> CFG -> ShowS
[CFG] -> ShowS
CFG -> String
(Int -> CFG -> ShowS)
-> (CFG -> String) -> ([CFG] -> ShowS) -> Show CFG
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CFG -> ShowS
showsPrec :: Int -> CFG -> ShowS
$cshow :: CFG -> String
show :: CFG -> String
$cshowList :: [CFG] -> ShowS
showList :: [CFG] -> ShowS
Show,(forall x. CFG -> Rep CFG x)
-> (forall x. Rep CFG x -> CFG) -> Generic CFG
forall x. Rep CFG x -> CFG
forall x. CFG -> Rep CFG x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CFG -> Rep CFG x
from :: forall x. CFG -> Rep CFG x
$cto :: forall x. Rep CFG x -> CFG
to :: forall x. Rep CFG x -> CFG
Generic,CFG -> CFG -> Bool
(CFG -> CFG -> Bool) -> (CFG -> CFG -> Bool) -> Eq CFG
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CFG -> CFG -> Bool
== :: CFG -> CFG -> Bool
$c/= :: CFG -> CFG -> Bool
/= :: CFG -> CFG -> Bool
Eq)


num_of_instructions :: CFG -> Int
num_of_instructions = (Int -> Int -> Int) -> Int -> IntMap Int -> Int
forall a b. (a -> b -> b) -> b -> IntMap a -> b
IM.foldr Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 (IntMap Int -> Int) -> (CFG -> IntMap Int) -> CFG -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int] -> Int) -> IntMap [Int] -> IntMap Int
forall a b. (a -> b) -> IntMap a -> IntMap b
IM.map [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (IntMap [Int] -> IntMap Int)
-> (CFG -> IntMap [Int]) -> CFG -> IntMap Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFG -> IntMap [Int]
cfg_blocks

init_cfg :: a -> CFG
init_cfg a
a = CFG { cfg_blocks :: IntMap [Int]
cfg_blocks = Int -> [Int] -> IntMap [Int]
forall a. Int -> a -> IntMap a
IM.singleton Int
0 [a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a], cfg_edges :: IntMap IntSet
cfg_edges = IntMap IntSet
forall a. IntMap a
IM.empty, cfg_addr_to_blockID :: IntMap Int
cfg_addr_to_blockID = Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IM.singleton (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a) Int
0, cfg_fresh :: Int
cfg_fresh = Int
1, cfg_instrs :: IntMap [Instruction]
cfg_instrs = IntMap [Instruction]
forall a. IntMap a
IM.empty }



-- | Returns true if the given blockID is a leaf-node in the given CFG.
is_end_node ::
  CFG     -- ^ The CFG
  -> Int  -- ^ The blockID
  -> Bool
is_end_node :: CFG -> Int -> Bool
is_end_node CFG
cfg = IntSet -> Bool
IS.null (IntSet -> Bool) -> (Int -> IntSet) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFG -> Int -> IntSet
post CFG
cfg

-- | The set of next blocks from the given block in a CFG
post :: CFG -> Int -> IS.IntSet
post :: CFG -> Int -> IntSet
post CFG
g Int
blockId = IntSet -> Maybe IntSet -> IntSet
forall a. a -> Maybe a -> a
fromMaybe IntSet
IS.empty (Int -> IntMap IntSet -> Maybe IntSet
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
blockId (CFG -> IntMap IntSet
cfg_edges CFG
g))

-- | Fetching an instruction list given a block ID
fetch_block ::
  CFG    -- ^ The CFG
  -> Int -- ^ The blockID
  -> [Instruction]
fetch_block :: CFG -> Int -> [Instruction]
fetch_block CFG
g Int
blockId =
  case Int -> IntMap [Instruction] -> Maybe [Instruction]
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
blockId (IntMap [Instruction] -> Maybe [Instruction])
-> IntMap [Instruction] -> Maybe [Instruction]
forall a b. (a -> b) -> a -> b
$ CFG -> IntMap [Instruction]
cfg_instrs (CFG -> IntMap [Instruction]) -> CFG -> IntMap [Instruction]
forall a b. (a -> b) -> a -> b
$ CFG
g of
    Maybe [Instruction]
Nothing -> String -> [Instruction]
forall a. HasCallStack => String -> a
error (String -> [Instruction]) -> String -> [Instruction]
forall a b. (a -> b) -> a -> b
$ String
"Block with ID" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
blockId String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" not found in cfg."
    Just [Instruction]
b -> [Instruction]
b


instance Cereal.Serialize CFG
instance NFData CFG