{-# LANGUAGE DuplicateRecordFields #-}

module Binja.ControlFlowGraph
  ( Binja.ControlFlowGraph.create,
    Binja.ControlFlowGraph.blocks,
    Binja.ControlFlowGraph.order,
    Binja.ControlFlowGraph.size,
  )
where

import Binja.BasicBlock
import Binja.Types (BNMlilSSAFunctionPtr, BasicBlockMlilSSA (..), CFGContext (..))
import Data.List (find)
import qualified Data.Map as Map
import qualified Data.Set as Set

-- outgoing edges becomes a lookup into the graph
--
-- Note: this could be done once during creation and cached
-- incoming edges becomces filter of the graph of keys with children containing block of interest
--

create :: BNMlilSSAFunctionPtr -> IO Binja.Types.CFGContext
create :: BNMlilSSAFunctionPtr -> IO CFGContext
create BNMlilSSAFunctionPtr
handle' = do
  -- blocks in function
  rawBlocks <- BNMlilSSAFunctionPtr -> IO [BNBasicBlockPtr]
Binja.BasicBlock.fromMlilSSAFunction BNMlilSSAFunctionPtr
handle'
  liftedBlocks <- mapM Binja.BasicBlock.fromBlockPtr rawBlocks
  -- entry block
  entryBlock' <-
    case Data.List.find ((0 ==) . start) liftedBlocks of
      Maybe BasicBlockMlilSSA
Nothing -> [Char] -> IO BasicBlockMlilSSA
forall a. HasCallStack => [Char] -> a
error [Char]
"Binja.ControlFlowGraph.create: No entry block found."
      Just BasicBlockMlilSSA
bb -> BasicBlockMlilSSA -> IO BasicBlockMlilSSA
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BasicBlockMlilSSA
bb
  -- edges from blocks
  rawOutgoingEdges <- mapM Binja.BasicBlock.outgoingEdges rawBlocks
  outgoingEdges' <- mapM (mapM Binja.BasicBlock.fromBlockEdge) rawOutgoingEdges
  let graph' =
        [(BasicBlockMlilSSA, Set BasicBlockEdge)]
-> Map BasicBlockMlilSSA (Set BasicBlockEdge)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(BasicBlockMlilSSA, Set BasicBlockEdge)]
 -> Map BasicBlockMlilSSA (Set BasicBlockEdge))
-> [(BasicBlockMlilSSA, Set BasicBlockEdge)]
-> Map BasicBlockMlilSSA (Set BasicBlockEdge)
forall a b. (a -> b) -> a -> b
$
          (BasicBlockMlilSSA
 -> [BasicBlockEdge] -> (BasicBlockMlilSSA, Set BasicBlockEdge))
-> [BasicBlockMlilSSA]
-> [[BasicBlockEdge]]
-> [(BasicBlockMlilSSA, Set BasicBlockEdge)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\BasicBlockMlilSSA
vertex [BasicBlockEdge]
edge -> (BasicBlockMlilSSA
vertex, [BasicBlockEdge] -> Set BasicBlockEdge
forall a. Ord a => [a] -> Set a
Set.fromList [BasicBlockEdge]
edge)) [BasicBlockMlilSSA]
liftedBlocks [[BasicBlockEdge]]
outgoingEdges'
  pure $ Binja.Types.CFGContext {entry = entryBlock', graph = graph'}

-- | List of blocks making up function
blocks :: Binja.Types.CFGContext -> [BasicBlockMlilSSA]
blocks :: CFGContext -> [BasicBlockMlilSSA]
blocks = Map BasicBlockMlilSSA (Set BasicBlockEdge) -> [BasicBlockMlilSSA]
forall k a. Map k a -> [k]
Map.keys (Map BasicBlockMlilSSA (Set BasicBlockEdge) -> [BasicBlockMlilSSA])
-> (CFGContext -> Map BasicBlockMlilSSA (Set BasicBlockEdge))
-> CFGContext
-> [BasicBlockMlilSSA]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFGContext -> Map BasicBlockMlilSSA (Set BasicBlockEdge)
graph

-- | Number of nodes
order :: Binja.Types.CFGContext -> Int
order :: CFGContext -> Int
order = Map BasicBlockMlilSSA (Set BasicBlockEdge) -> Int
forall k a. Map k a -> Int
Map.size (Map BasicBlockMlilSSA (Set BasicBlockEdge) -> Int)
-> (CFGContext -> Map BasicBlockMlilSSA (Set BasicBlockEdge))
-> CFGContext
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFGContext -> Map BasicBlockMlilSSA (Set BasicBlockEdge)
graph

-- | Numer of edges
size :: Binja.Types.CFGContext -> Int
size :: CFGContext -> Int
size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> (CFGContext -> [Int]) -> CFGContext -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set BasicBlockEdge -> Int) -> [Set BasicBlockEdge] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Set BasicBlockEdge -> Int
forall a. Set a -> Int
Set.size ([Set BasicBlockEdge] -> [Int])
-> (CFGContext -> [Set BasicBlockEdge]) -> CFGContext -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map BasicBlockMlilSSA (Set BasicBlockEdge) -> [Set BasicBlockEdge]
forall k a. Map k a -> [a]
Map.elems (Map BasicBlockMlilSSA (Set BasicBlockEdge)
 -> [Set BasicBlockEdge])
-> (CFGContext -> Map BasicBlockMlilSSA (Set BasicBlockEdge))
-> CFGContext
-> [Set BasicBlockEdge]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFGContext -> Map BasicBlockMlilSSA (Set BasicBlockEdge)
graph