content: {-# LANGUAGE Strict #-}
  
  {- |
    Module      :  GHC.CmmToAsm.CFG.Dominators
    Copyright   :  (c) Matt Morrow 2009
    License     :  BSD3
    Maintainer  :  <klebinger.andreas@gmx.at>
    Stability   :  stable
    Portability :  portable
  
    The Lengauer-Tarjan graph dominators algorithm.
  
      \[1\] Lengauer, Tarjan,
        /A Fast Algorithm for Finding Dominators in a Flowgraph/, 1979.
  
      \[2\] Muchnick,
        /Advanced Compiler Design and Implementation/, 1997.
  
      \[3\] Brisk, Sarrafzadeh,
        /Interference CGraphs for Procedures in Static Single/
        /Information Form are Interval CGraphs/, 2007.
  
   * Strictness
  
   Unless stated otherwise all exposed functions might fully evaluate their input
   but are not guaranteed to do so.
  
  -}
  
  module GHC.CmmToAsm.CFG.Dominators (
     Node,Path,Edge
    ,Graph,Rooted
    ,idom,ipdom
    ,domTree,pdomTree
    ,dom,pdom
    ,pddfs,rpddfs
    ,fromAdj,fromEdges
    ,toAdj,toEdges
    ,asTree,asCGraph
    ,parents,ancestors
  ) where
  
  import GHC.Prelude
  import Data.Bifunctor
  import Data.Tuple (swap)
  
  import Data.Tree
  import Data.IntMap(IntMap)
  import Data.IntSet(IntSet)
  import qualified Data.IntMap.Strict as IM
  import qualified Data.IntSet as IS
  
  import Control.Monad
  import Control.Monad.ST.Strict
  
  import Data.Array.ST
  import Data.Array.Base
    (unsafeNewArray_
    ,unsafeWrite,unsafeRead)
  import GHC.Data.Word64Set (Word64Set)
  import qualified GHC.Data.Word64Set as WS
  import GHC.Data.Word64Map (Word64Map)
  import qualified GHC.Data.Word64Map as WM
  import Data.Word
  
  -----------------------------------------------------------------------------
  
  -- Compacted nodes; these can be stored in contiguous arrays
  type CNode       = Int
  type CGraph      = IntMap IntSet
  
  type Node     = Word64
  type Path     = [Node]
  type Edge     = (Node, Node)
  type Graph    = Word64Map Word64Set
  type Rooted   = (Node, Graph)
  
  -----------------------------------------------------------------------------
  
  -- | /Dominators/.
  -- Complexity as for @idom@
  dom :: Rooted -> [(Node, Path)]
  dom = ancestors . domTree
  
  -- | /Post-dominators/.
  -- Complexity as for @idom@.
  pdom :: Rooted -> [(Node, Path)]
  pdom = ancestors . pdomTree
  
  -- | /Dominator tree/.
  -- Complexity as for @idom@.
  domTree :: Rooted -> Tree Node
  domTree a@(r,_) =
    let is = filter ((/=r).fst) (idom a)
        tg = fromEdges (fmap swap is)
    in asTree (r,tg)
  
  -- | /Post-dominator tree/.
  -- Complexity as for @idom@.
  pdomTree :: Rooted -> Tree Node
  pdomTree a@(r,_) =
    let is = filter ((/=r).fst) (ipdom a)
        tg = fromEdges (fmap swap is)
    in asTree (r,tg)
  
  -- | /Immediate dominators/.
  -- /O(|E|*alpha(|E|,|V|))/, where /alpha(m,n)/ is
  -- \"a functional inverse of Ackermann's function\".
  --
  -- This Complexity bound assumes /O(1)/ indexing. Since we're
  -- using @IntMap@, it has an additional /lg |V|/ factor
  -- somewhere in there. I'm not sure where.
  idom :: Rooted -> [(Node,Node)]
  idom rg = runST (evalS idomM =<< initEnv (pruneReach rg))
  
  -- | /Immediate post-dominators/.
  -- Complexity as for @idom@.
  ipdom :: Rooted -> [(Node,Node)]
  ipdom rg = runST (evalS idomM =<< initEnv (pruneReach (second predGW rg)))
  
  -----------------------------------------------------------------------------
  
  -- | /Post-dominated depth-first search/.
  pddfs :: Rooted -> [Node]
  pddfs = reverse . rpddfs
  
  -- | /Reverse post-dominated depth-first search/.
  rpddfs :: Rooted -> [Node]
  rpddfs = concat . levels . pdomTree
  
  -----------------------------------------------------------------------------
  
  type Dom s a = S s (Env s) a
  type NodeSet    = Word64Set
  type NodeMap a  = Word64Map a
  data Env s = Env
    {succE      :: !CGraph
    ,predE      :: !CGraph
    ,bucketE    :: !CGraph
    ,dfsE       :: {-# UNPACK #-}!Int
    ,zeroE      :: {-# UNPACK #-}!CNode
    ,rootE      :: {-# UNPACK #-}!CNode
    ,labelE     :: {-# UNPACK #-}!(Arr s CNode)
    ,parentE    :: {-# UNPACK #-}!(Arr s CNode)
    ,ancestorE  :: {-# UNPACK #-}!(Arr s CNode)
    ,childE     :: {-# UNPACK #-}!(Arr s CNode)
    ,ndfsE      :: {-# UNPACK #-}!(Arr s CNode)
    ,dfnE       :: {-# UNPACK #-}!(Arr s Int)
    ,sdnoE      :: {-# UNPACK #-}!(Arr s Int)
    ,sizeE      :: {-# UNPACK #-}!(Arr s Int)
    ,domE       :: {-# UNPACK #-}!(Arr s CNode)
    ,rnE        :: {-# UNPACK #-}!(Arr s Node)}
  
  -----------------------------------------------------------------------------
  
  idomM :: Dom s [(Node,Node)]
  idomM = do
    dfsDom =<< rootM
    n <- gets dfsE
    forM_ [n,n-1..1] (\i-> do
      w <- ndfsM i
      ps <- predsM w
      forM_ ps (\v-> do
        sw <- sdnoM w
        u <- eval v
        su <- sdnoM u
        when (su < sw)
          (store sdnoE w su))
      z <- ndfsM =<< sdnoM w
      modify(\e->e{bucketE=IM.adjust
                        (w`IS.insert`)
                        z (bucketE e)})
      pw <- parentM w
      link pw w
      bps <- bucketM pw
      forM_ bps (\v-> do
        u <- eval v
        su <- sdnoM u
        sv <- sdnoM v
        let dv = case su < sv of
                  True-> u
                  False-> pw
        store domE v dv))
    forM_ [1..n] (\i-> do
      w <- ndfsM i
      j <- sdnoM w
      z <- ndfsM j
      dw <- domM w
      when (dw /= z)
        (do ddw <- domM dw
            store domE w ddw))
    fromEnv
  
  -----------------------------------------------------------------------------
  
  eval :: CNode -> Dom s CNode
  eval v = do
    n0 <- zeroM
    a  <- ancestorM v
    case a==n0 of
      True-> labelM v
      False-> do
        compress v
        a   <- ancestorM v
        l   <- labelM v
        la  <- labelM a
        sl  <- sdnoM l
        sla <- sdnoM la
        case sl <= sla of
          True-> return l
          False-> return la
  
  compress :: CNode -> Dom s ()
  compress v = do
    n0  <- zeroM
    a   <- ancestorM v
    aa  <- ancestorM a
    when (aa /= n0) (do
      compress a
      a   <- ancestorM v
      aa  <- ancestorM a
      l   <- labelM v
      la  <- labelM a
      sl  <- sdnoM l
      sla <- sdnoM la
      when (sla < sl)
        (store labelE v la)
      store ancestorE v aa)
  
  -----------------------------------------------------------------------------
  
  link :: CNode -> CNode -> Dom s ()
  link v w = do
    n0  <- zeroM
    lw  <- labelM w
    slw <- sdnoM lw
    let balance s = do
          c   <- childM s
          lc  <- labelM c
          slc <- sdnoM lc
          case slw < slc of
            False-> return s
            True-> do
              zs  <- sizeM s
              zc  <- sizeM c
              cc  <- childM c
              zcc <- sizeM cc
              case 2*zc <= zs+zcc of
                True-> do
                  store ancestorE c s
                  store childE s cc
                  balance s
                False-> do
                  store sizeE c zs
                  store ancestorE s c
                  balance c
    s   <- balance w
    lw  <- labelM w
    zw  <- sizeM w
    store labelE s lw
    store sizeE v . (+zw) =<< sizeM v
    let follow s =
          when (s /= n0) (do
            store ancestorE s v
            follow =<< childM s)
    zv  <- sizeM v
    follow =<< case zv < 2*zw of
                False-> return s
                True-> do
                  cv <- childM v
                  store childE v s
                  return cv
  
  -----------------------------------------------------------------------------
  
  dfsDom :: CNode -> Dom s ()
  dfsDom i = do
    _   <- go i
    n0  <- zeroM
    r   <- rootM
    store parentE r n0
    where go i = do
            n <- nextM
            store dfnE   i n
            store sdnoE  i n
            store ndfsE  n i
            store labelE i i
            ss <- succsM i
            forM_ ss (\j-> do
              s <- sdnoM j
              case s==0 of
                False-> return()
                True-> do
                  store parentE j i
                  go j)
  
  -----------------------------------------------------------------------------
  
  initEnv :: Rooted -> ST s (Env s)
  initEnv (r0,g0) = do
    -- CGraph renumbered to indices from 1 to |V|
    let (g,rnmap) = renum 1 g0
        pred      = predG g -- reverse graph
        root      = rnmap WM.! r0 -- renamed root
        n         = IM.size g
        ns        = [0..n]
        m         = n+1
  
    let bucket = IM.fromList
          (zip ns (repeat mempty))
  
    rna <- newW m
    writes rna (fmap swap
          (WM.toList rnmap))
  
    doms      <- newI m
    sdno      <- newI m
    size      <- newI m
    parent    <- newI m
    ancestor  <- newI m
    child     <- newI m
    label     <- newI m
    ndfs      <- newI m
    dfn       <- newI m
  
    -- Initialize all arrays
    forM_ [0..n] (doms.=0)
    forM_ [0..n] (sdno.=0)
    forM_ [1..n] (size.=1)
    forM_ [0..n] (ancestor.=0)
    forM_ [0..n] (child.=0)
  
    (doms.=root) root
    (size.=0) 0
    (label.=0) 0
  
    return (Env
      {rnE        = rna
      ,dfsE       = 0
      ,zeroE      = 0
      ,rootE      = root
      ,labelE     = label
      ,parentE    = parent
      ,ancestorE  = ancestor
      ,childE     = child
      ,ndfsE      = ndfs
      ,dfnE       = dfn
      ,sdnoE      = sdno
      ,sizeE      = size
      ,succE      = g
      ,predE      = pred
      ,bucketE    = bucket
      ,domE       = doms})
  
  fromEnv :: Dom s [(Node,Node)]
  fromEnv = do
    dom   <- gets domE
    rn    <- gets rnE
    -- r     <- gets rootE
    (_,n) <- st (getBounds dom)
    forM [1..n] (\i-> do
      j <- st(rn!:i)
      d <- st(dom!:i)
      k <- st(rn!:d)
      return (j,k))
  
  -----------------------------------------------------------------------------
  
  zeroM :: Dom s CNode
  zeroM = gets zeroE
  domM :: CNode -> Dom s CNode
  domM = fetch domE
  rootM :: Dom s CNode
  rootM = gets rootE
  succsM :: CNode -> Dom s [CNode]
  succsM i = gets (IS.toList . (! i) . succE)
  predsM :: CNode -> Dom s [CNode]
  predsM i = gets (IS.toList . (! i) . predE)
  bucketM :: CNode -> Dom s [CNode]
  bucketM i = gets (IS.toList . (! i) . bucketE)
  sizeM :: CNode -> Dom s Int
  sizeM = fetch sizeE
  sdnoM :: CNode -> Dom s Int
  sdnoM = fetch sdnoE
  -- dfnM :: CNode -> Dom s Int
  -- dfnM = fetch dfnE
  ndfsM :: Int -> Dom s CNode
  ndfsM = fetch ndfsE
  childM :: CNode -> Dom s CNode
  childM = fetch childE
  ancestorM :: CNode -> Dom s CNode
  ancestorM = fetch ancestorE
  parentM :: CNode -> Dom s CNode
  parentM = fetch parentE
  labelM :: CNode -> Dom s CNode
  labelM = fetch labelE
  nextM :: Dom s Int
  nextM = do
    n <- gets dfsE
    let n' = n+1
    modify(\e->e{dfsE=n'})
    return n'
  
  -----------------------------------------------------------------------------
  
  type A = STUArray
  type Arr s a = A s Int a
  
  infixl 9 !:
  infixr 2 .=
  
  -- | arr .= x idx => write x to index
  (.=) :: (MArray (A s) a (ST s))
       => Arr s a -> a -> Int -> ST s ()
  (v .= x) i = unsafeWrite v i x
  
  (!:) :: (MArray (A s) a (ST s))
       => A s Int a -> Int -> ST s a
  a !: i = do
    o <- unsafeRead a i
    return $! o
  
  new :: (MArray (A s) a (ST s))
      => Int -> ST s (Arr s a)
  new n = unsafeNewArray_ (0,n-1)
  
  newI :: Int -> ST s (Arr s Int)
  newI = new
  
  newW :: Int -> ST s (Arr s Node)
  newW = new
  
  writes :: (MArray (A s) a (ST s))
       => Arr s a -> [(Int,a)] -> ST s ()
  writes a xs = forM_ xs (\(i,x) -> (a.=x) i)
  
  
  (!) :: Monoid a => IntMap a -> Int -> a
  (!) g n = maybe mempty id (IM.lookup n g)
  
  fromAdj :: [(Node, [Node])] -> Graph
  fromAdj = WM.fromList . fmap (second WS.fromList)
  
  fromEdges :: [Edge] -> Graph
  fromEdges = collectW WS.union fst (WS.singleton . snd)
  
  toAdj :: Graph -> [(Node, [Node])]
  toAdj = fmap (second WS.toList) . WM.toList
  
  toEdges :: Graph -> [Edge]
  toEdges = concatMap (uncurry (fmap . (,))) . toAdj
  
  predG :: CGraph -> CGraph
  predG g = IM.unionWith IS.union (go g) g0
    where g0 = fmap (const mempty) g
          go = flip IM.foldrWithKey mempty (\i a m ->
                  foldl' (\m p -> IM.insertWith mappend p
                                        (IS.singleton i) m)
                          m
                         (IS.toList a))
  
  predGW :: Graph -> Graph
  predGW g = WM.unionWith WS.union (go g) g0
    where g0 = fmap (const mempty) g
          go = flip WM.foldrWithKey mempty (\i a m ->
                  foldl' (\m p -> WM.insertWith mappend p
                                        (WS.singleton i) m)
                          m
                         (WS.toList a))
  
  pruneReach :: Rooted -> Rooted
  pruneReach (r,g) = (r,g2)
    where is = reachable
                (maybe mempty id
                  . flip WM.lookup g) $ r
          g2 = WM.fromList
              . fmap (second (WS.filter (`WS.member`is)))
              . filter ((`WS.member`is) . fst)
              . WM.toList $ g
  
  tip :: Tree a -> (a, [Tree a])
  tip (Node a ts) = (a, ts)
  
  parents :: Tree a -> [(a, a)]
  parents (Node i xs) = p i xs
          ++ concatMap parents xs
    where p i = fmap (flip (,) i . rootLabel)
  
  ancestors :: Tree a -> [(a, [a])]
  ancestors = go []
    where go acc (Node i xs)
            = let acc' = i:acc
              in p acc' xs ++ concatMap (go acc') xs
          p is = fmap (flip (,) is . rootLabel)
  
  asCGraph :: Tree Node -> Rooted
  asCGraph t@(Node a _) = let g = go t in (a, fromAdj g)
    where go (Node a ts) = let as = (fst . unzip . fmap tip) ts
                            in (a, as) : concatMap go ts
  
  asTree :: Rooted -> Tree Node
  asTree (r,g) = let go a = Node a (fmap go ((WS.toList . f) a))
                     f = (g !)
              in go r
    where (!) g n = maybe mempty id (WM.lookup n g)
  
  
  reachable :: (Node -> NodeSet) -> (Node -> NodeSet)
  reachable f a = go (WS.singleton a) a
    where go seen a = let s = f a
                          as = WS.toList (s `WS.difference` seen)
                      in foldl' go (s `WS.union` seen) as
  
  collectW :: (c -> c -> c)
          -> (a -> Node) -> (a -> c) -> [a] -> Word64Map c
  collectW (<>) f g
    = foldl' (\m a -> WM.insertWith (<>)
                                    (f a)
                                    (g a) m) mempty
  
  -- | renum n g: Rename all nodes
  --
  -- Gives nodes sequential names starting at n.
  -- Returns the new graph and a mapping.
  -- (renamed, old -> new)
  renum :: Int -> Graph -> (CGraph, NodeMap CNode)
  renum from = (\(_,m,g)->(g,m))
    . WM.foldrWithKey
        (\i ss (!n,!env,!new)->
            let (j,n2,env2) = go n env i
                (n3,env3,ss2) = WS.fold
                  (\k (!n,!env,!new)->
                      case go n env k of
                        (l,n2,env2)-> (n2,env2,l `IS.insert` new))
                  (n2,env2,mempty) ss
                new2 = IM.insertWith IS.union j ss2 new
            in (n3,env3,new2)) (from,mempty,mempty)
    where go :: Int
             -> NodeMap CNode
             -> Node
             -> (CNode,Int,NodeMap CNode)
          go !n !env i =
            case WM.lookup i env of
              Just j -> (j,n,env)
              Nothing -> (n,n+1,WM.insert i n env)
  
  -----------------------------------------------------------------------------
  
  -- Nothing better than reinventing the state monad.
  newtype S z s a = S {unS :: forall o. (a -> s -> ST z o) -> s -> ST z o}
    deriving (Functor)
  instance Monad (S z s) where
    return = pure
    S g >>= f = S (\k -> g (\a -> unS (f a) k))
  instance Applicative (S z s) where
    pure a = S (\k -> k a)
    (<*>) = ap
  -- get :: S z s s
  -- get = S (\k s -> k s s)
  gets :: (s -> a) -> S z s a
  gets f = S (\k s -> k (f s) s)
  -- set :: s -> S z s ()
  -- set s = S (\k _ -> k () s)
  modify :: (s -> s) -> S z s ()
  modify f = S (\k -> k () . f)
  -- runS :: S z s a -> s -> ST z (a, s)
  -- runS (S g) = g (\a s -> return (a,s))
  evalS :: S z s a -> s -> ST z a
  evalS (S g) = g ((return .) . const)
  -- execS :: S z s a -> s -> ST z s
  -- execS (S g) = g ((return .) . flip const)
  st :: ST z a -> S z s a
  st m = S (\k s-> do
    a <- m
    k a s)
  store :: (MArray (A z) a (ST z))
        => (s -> Arr z a) -> Int -> a -> S z s ()
  store f i x = do
    a <- gets f
    st ((a.=x) i)
  fetch :: (MArray (A z) a (ST z))
        => (s -> Arr z a) -> Int -> S z s a
  fetch f i = do
    a <- gets f
    st (a!:i)
  