-- Find Dominators of a graph.
--
-- Author: Bertram Felgenhauer <int-e@gmx.de>
--
-- Implementation based on
-- Keith D. Cooper, Timothy J. Harvey, Ken Kennedy,
-- "A Simple, Fast Dominance Algorithm",
-- (http://citeseer.ist.psu.edu/cooper01simple.html)

module Data.Graph.Inductive.Query.Dominators (
    dom,
    iDom
) where

import           Data.Array
import           Data.Graph.Inductive.Graph
import           Data.Graph.Inductive.Query.DFS
import           Data.IntMap                    (IntMap)
import qualified Data.IntMap                    as I
import           Data.Maybe (mapMaybe)
import           Data.Tree                      (Tree (..))
import qualified Data.Tree                      as T

{-# ANN iDom "HLint: ignore Use ***" #-}
-- | return immediate dominators for each reachable node of a graph, given a root
iDom :: (Graph gr) => gr a b -> Node -> [(Node,Node)]
iDom :: forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> [(Int, Int)]
iDom gr a b
g Int
root = let (IDom
result, IDom
toNode, FromNode
_) = gr a b -> Int -> (IDom, IDom, FromNode)
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> (IDom, IDom, FromNode)
idomWork gr a b
g Int
root
              in  ((Int, Int) -> (Int, Int)) -> [(Int, Int)] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Int
a, Int
b) -> (IDom
toNode IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
a, IDom
toNode IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
b)) (IDom -> [(Int, Int)]
forall i e. Ix i => Array i e -> [(i, e)]
assocs IDom
result)

-- | return the set of dominators of the reachable nodes of a graph, given a root
dom :: (Graph gr) => gr a b -> Node -> [(Node,[Node])]
dom :: forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> [(Int, [Int])]
dom gr a b
g Int
root = let
    (IDom
iD, IDom
toNode, FromNode
fromNode) = gr a b -> Int -> (IDom, IDom, FromNode)
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> (IDom, IDom, FromNode)
idomWork gr a b
g Int
root
    dom' :: Array Int [Int]
dom' = IDom -> IDom -> Array Int [Int]
getDom IDom
toNode IDom
iD
    nodes' :: [Int]
nodes' = gr a b -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [Int]
nodes gr a b
g
    rest :: [Int]
rest = FromNode -> [Int]
forall a. IntMap a -> [Int]
I.keys ((Int -> Bool) -> FromNode -> FromNode
forall a. (a -> Bool) -> IntMap a -> IntMap a
I.filter (-Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==) FromNode
fromNode)
  in
    [(IDom
toNode IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
i, Array Int [Int]
dom' Array Int [Int] -> Int -> [Int]
forall i e. Ix i => Array i e -> i -> e
! Int
i) | Int
i <- (Int, Int) -> [Int]
forall a. Ix a => (a, a) -> [a]
range (Array Int [Int] -> (Int, Int)
forall i e. Array i e -> (i, i)
bounds Array Int [Int]
dom')] [(Int, [Int])] -> [(Int, [Int])] -> [(Int, [Int])]
forall a. [a] -> [a] -> [a]
++
    [(Int
n, [Int]
nodes') | Int
n <- [Int]
rest]

-- internal node type
type Node' = Int
-- array containing the immediate dominator of each node, or an approximation
-- thereof. the dominance set of a node can be found by taking the union of
-- {node} and the dominance set of its immediate dominator.
type IDom = Array Node' Node'
-- array containing the list of predecessors of each node
type Preds = Array Node' [Node']
-- arrays for translating internal nodes back to graph nodes and back
type ToNode = Array Node' Node
type FromNode = IntMap Node'

idomWork :: (Graph gr) => gr a b -> Node -> (IDom, ToNode, FromNode)
idomWork :: forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> (IDom, IDom, FromNode)
idomWork gr a b
g Int
root =
  case [Int] -> gr a b -> [Tree Int]
forall (gr :: * -> * -> *) a b.
Graph gr =>
[Int] -> gr a b -> [Tree Int]
dff [Int
root] gr a b
g of
    [] -> [Char] -> (IDom, IDom, FromNode)
forall a. HasCallStack => [Char] -> a
error [Char]
"Dominators.idomWork: root not in graph"
    Tree Int
tree : [Tree Int]
_ ->
      let
        nds :: [Int]
nds = Int -> gr a b -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> [Int]
reachable Int
root gr a b
g
        -- use depth first tree from root do build the first approximation
        -- relabel the tree so that paths from the root have increasing nodes
        (Int
s, Tree Int
ntree) = Int -> Tree Int -> (Int, Tree Int)
forall a. Int -> Tree a -> (Int, Tree Int)
numberTree Int
0 Tree Int
tree
        -- the approximation iDom0 just maps each node to its parent
        iD0 :: IDom
iD0 = (Int, Int) -> [(Int, Int)] -> IDom
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int
1, Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ([(Int, Int)] -> [(Int, Int)]
forall a. HasCallStack => [a] -> [a]
tail ([(Int, Int)] -> [(Int, Int)]) -> [(Int, Int)] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ Int -> Tree Int -> [(Int, Int)]
forall a. a -> Tree a -> [(a, a)]
treeEdges (-Int
1) Tree Int
ntree)
        -- fromNode translates graph nodes to relabeled (internal) nodes
        fromNode :: FromNode
fromNode = (Int -> Int -> Int) -> FromNode -> FromNode -> FromNode
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
I.unionWith Int -> Int -> Int
forall a b. a -> b -> a
const ([(Int, Int)] -> FromNode
forall a. [(Int, a)] -> IntMap a
I.fromList ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Tree Int -> [Int]
forall a. Tree a -> [a]
T.flatten Tree Int
tree) (Tree Int -> [Int]
forall a. Tree a -> [a]
T.flatten Tree Int
ntree))) ([(Int, Int)] -> FromNode
forall a. [(Int, a)] -> IntMap a
I.fromList ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
nds (Int -> [Int]
forall a. a -> [a]
repeat (-Int
1))))
        -- toNode translates internal nodes to graph nodes
        toNode :: IDom
toNode = (Int, Int) -> [(Int, Int)] -> IDom
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int
0, Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Tree Int -> [Int]
forall a. Tree a -> [a]
T.flatten Tree Int
ntree) (Tree Int -> [Int]
forall a. Tree a -> [a]
T.flatten Tree Int
tree))
        preds :: Array Int [Int]
preds = (Int, Int) -> [(Int, [Int])] -> Array Int [Int]
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int
1, Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [(Int
i, (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1) ((Int -> Maybe Int) -> [Int] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Int -> FromNode -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
`I.lookup` FromNode
fromNode)
                                (gr a b -> Int -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> Int -> [Int]
pre gr a b
g (IDom
toNode IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
i)))) | Int
i <- [Int
1..Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]]
        -- iteratively improve the approximation to find iDom.
        iD :: IDom
iD = (IDom -> IDom) -> IDom -> IDom
forall a. Eq a => (a -> a) -> a -> a
fixEq (Array Int [Int] -> IDom -> IDom
refineIDom Array Int [Int]
preds) IDom
iD0
      in (IDom
iD, IDom
toNode, FromNode
fromNode)

-- for each node in iDom, find the intersection of all its predecessor's
-- dominating sets, and update iDom accordingly.
refineIDom :: Preds -> IDom -> IDom
refineIDom :: Array Int [Int] -> IDom -> IDom
refineIDom Array Int [Int]
preds IDom
iD = ([Int] -> Int) -> Array Int [Int] -> IDom
forall a b. (a -> b) -> Array Int a -> Array Int b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> Int -> Int) -> [Int] -> Int
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 (IDom -> Int -> Int -> Int
intersect IDom
iD)) Array Int [Int]
preds

-- find the intersection of the two given dominance sets.
intersect :: IDom -> Node' -> Node' -> Node'
intersect :: IDom -> Int -> Int -> Int
intersect IDom
iD Int
a Int
b = case Int
a Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
b of
    Ordering
LT -> IDom -> Int -> Int -> Int
intersect IDom
iD Int
a (IDom
iD IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
b)
    Ordering
EQ -> Int
a
    Ordering
GT -> IDom -> Int -> Int -> Int
intersect IDom
iD (IDom
iD IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
a) Int
b

-- convert an IDom to dominance sets. we translate to graph nodes here
-- because mapping later would be more expensive and lose sharing.
getDom :: ToNode -> IDom -> Array Node' [Node]
getDom :: IDom -> IDom -> Array Int [Int]
getDom IDom
toNode IDom
iD = let
    res :: Array Int [Int]
res = (Int, Int) -> [(Int, [Int])] -> Array Int [Int]
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int
0, (Int, Int) -> Int
forall a b. (a, b) -> b
snd (IDom -> (Int, Int)
forall i e. Array i e -> (i, i)
bounds IDom
iD)) ((Int
0, [IDom
toNode IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
0]) (Int, [Int]) -> [(Int, [Int])] -> [(Int, [Int])]
forall a. a -> [a] -> [a]
:
          [(Int
i, IDom
toNode IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Array Int [Int]
res Array Int [Int] -> Int -> [Int]
forall i e. Ix i => Array i e -> i -> e
! (IDom
iD IDom -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
i)) | Int
i <- (Int, Int) -> [Int]
forall a. Ix a => (a, a) -> [a]
range (IDom -> (Int, Int)
forall i e. Array i e -> (i, i)
bounds IDom
iD)])
  in
    Array Int [Int]
res

-- relabel tree, labeling vertices with consecutive numbers in depth first order
numberTree :: Node' -> Tree a -> (Node', Tree Node')
numberTree :: forall a. Int -> Tree a -> (Int, Tree Int)
numberTree Int
n (Node a
_ [Tree a]
ts) = let (Int
n', [Tree Int]
ts') = Int -> [Tree a] -> (Int, [Tree Int])
forall a. Int -> [Tree a] -> (Int, [Tree Int])
numberForest (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Tree a]
ts
                           in  (Int
n', Int -> [Tree Int] -> Tree Int
forall a. a -> [Tree a] -> Tree a
Node Int
n [Tree Int]
ts')

-- same as numberTree, for forests.
numberForest :: Node' -> [Tree a] -> (Node', [Tree Node'])
numberForest :: forall a. Int -> [Tree a] -> (Int, [Tree Int])
numberForest Int
n []     = (Int
n, [])
numberForest Int
n (Tree a
t:[Tree a]
ts) = let (Int
n', Tree Int
t')   = Int -> Tree a -> (Int, Tree Int)
forall a. Int -> Tree a -> (Int, Tree Int)
numberTree Int
n Tree a
t
                            (Int
n'', [Tree Int]
ts') = Int -> [Tree a] -> (Int, [Tree Int])
forall a. Int -> [Tree a] -> (Int, [Tree Int])
numberForest Int
n' [Tree a]
ts
                        in  (Int
n'', Tree Int
t'Tree Int -> [Tree Int] -> [Tree Int]
forall a. a -> [a] -> [a]
:[Tree Int]
ts')

-- return the edges of the tree, with an added dummy root node.
treeEdges :: a -> Tree a -> [(a,a)]
treeEdges :: forall a. a -> Tree a -> [(a, a)]
treeEdges a
a (Node a
b [Tree a]
ts) = (a
b,a
a) (a, a) -> [(a, a)] -> [(a, a)]
forall a. a -> [a] -> [a]
: (Tree a -> [(a, a)]) -> [Tree a] -> [(a, a)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (a -> Tree a -> [(a, a)]
forall a. a -> Tree a -> [(a, a)]
treeEdges a
b) [Tree a]
ts

-- find a fixed point of f, iteratively
fixEq :: (Eq a) => (a -> a) -> a -> a
fixEq :: forall a. Eq a => (a -> a) -> a -> a
fixEq a -> a
f a
v | a
v' a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
v   = a
v
          | Bool
otherwise = (a -> a) -> a -> a
forall a. Eq a => (a -> a) -> a -> a
fixEq a -> a
f a
v'
    where v' :: a
v' = a -> a
f a
v

{-
:m +Data.Graph.Inductive
let g0 = mkGraph [(i,()) | i <- [0..4]] [(a,b,()) | (a,b) <- [(0,1),(1,2),(0,3),(3,2),(4,0)]] :: Gr () ()
let g1 = mkGraph [(i,()) | i <- [0..4]] [(a,b,()) | (a,b) <- [(0,1),(1,2),(2,3),(1,3),(3,4)]] :: Gr () ()
let g2,g3,g4 :: Int -> Gr () (); g2 n = mkGraph [(i,()) | i <- [0..n-1]] ([(a,a+1,()) | a <- [0..n-2]] ++ [(a,a+2,()) | a <- [0..n-3]]); g3 n =mkGraph [(i,()) | i <- [0..n-1]] ([(a,a+2,()) | a <- [0..n-3]] ++ [(a,a+1,()) | a <- [0..n-2]]); g4 n =mkGraph [(i,()) | i <- [0..n-1]] ([(a+2,a,()) | a <- [0..n-3]] ++ [(a+1,a,()) | a <- [0..n-2]])
:m -Data.Graph.Inductive
-}