-- | An abstract interface for a fast reachability data structure constructed
-- from a 'GHC.Data.Graph.Directed' graph.
module GHC.Data.Graph.Directed.Reachability
  ( ReachabilityIndex

  -- * Constructing a reachability index
  , graphReachability, cyclicGraphReachability

  -- * Reachability queries
  , allReachable, allReachableMany
  , isReachable, isReachableMany
  )
  where

import GHC.Prelude
import GHC.Data.Maybe

import qualified Data.Graph as G
import Data.Graph ( Vertex, SCC(..) )

import Data.Array ((!))
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS

import GHC.Data.Graph.Directed.Internal

--------------------------------------------------------------------------------
-- * Reachability index
--------------------------------------------------------------------------------

-- | The abstract data structure for fast reachability queries
data ReachabilityIndex node = ReachabilityIndex {
    forall node. ReachabilityIndex node -> IntMap IntSet
index :: IM.IntMap IS.IntSet,
    forall node. ReachabilityIndex node -> Vertex -> node
from_vertex :: Vertex -> node,
    forall node. ReachabilityIndex node -> node -> Maybe Vertex
to_vertex :: node -> Maybe Vertex
}

--------------------------------------------------------------------------------
-- * Construction
--------------------------------------------------------------------------------

-- | Construct a 'ReachabilityIndex' from an acyclic 'Graph'.
-- If the graph can have cycles, use 'cyclicGraphReachability'
graphReachability :: Graph node -> ReachabilityIndex node
graphReachability :: forall node. Graph node -> ReachabilityIndex node
graphReachability (Graph IntGraph
g Vertex -> node
from node -> Maybe Vertex
to) =
  ReachabilityIndex{index :: IntMap IntSet
index = IntMap IntSet
reachableGraph, from_vertex :: Vertex -> node
from_vertex = Vertex -> node
from, to_vertex :: node -> Maybe Vertex
to_vertex = node -> Maybe Vertex
to}
    where
      reachableGraph :: IM.IntMap IS.IntSet
      reachableGraph :: IntMap IntSet
reachableGraph = [(Vertex, IntSet)] -> IntMap IntSet
forall a. [(Vertex, a)] -> IntMap a
IM.fromList [(Vertex
v, Vertex -> IntSet
do_one Vertex
v) | Vertex
v <- IntGraph -> [Vertex]
G.vertices IntGraph
g]

      do_one :: Vertex -> IntSet
do_one Vertex
v = [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([Vertex] -> IntSet
IS.fromList (IntGraph
g IntGraph -> Vertex -> [Vertex]
forall i e. Ix i => Array i e -> i -> e
! Vertex
v) IntSet -> [IntSet] -> [IntSet]
forall a. a -> [a] -> [a]
: (Vertex -> Maybe IntSet) -> [Vertex] -> [IntSet]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((Vertex -> IntMap IntSet -> Maybe IntSet)
-> IntMap IntSet -> Vertex -> Maybe IntSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Vertex -> IntMap IntSet -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup IntMap IntSet
reachableGraph) (IntGraph
g IntGraph -> Vertex -> [Vertex]
forall i e. Ix i => Array i e -> i -> e
! Vertex
v))

-- | Construct a 'ReachabilityIndex' from a 'Graph' which may have cycles.
-- If this reachability index is just going to be used once, it may make sense
-- to use 'reachablesG' instead, which will traverse the reachable nodes without
-- constructing the index -- which may be faster.
cyclicGraphReachability :: Graph node -> ReachabilityIndex node
cyclicGraphReachability :: forall node. Graph node -> ReachabilityIndex node
cyclicGraphReachability (Graph IntGraph
g Vertex -> node
from node -> Maybe Vertex
to) =
  ReachabilityIndex{index :: IntMap IntSet
index = IntMap IntSet
reachableGraphCyclic, from_vertex :: Vertex -> node
from_vertex = Vertex -> node
from, to_vertex :: node -> Maybe Vertex
to_vertex = node -> Maybe Vertex
to}
    where
      reachableGraphCyclic :: IM.IntMap IS.IntSet
      reachableGraphCyclic :: IntMap IntSet
reachableGraphCyclic = (IntMap IntSet -> SCC Vertex -> IntMap IntSet)
-> IntMap IntSet -> [SCC Vertex] -> IntMap IntSet
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IntMap IntSet -> SCC Vertex -> IntMap IntSet
add_one_comp IntMap IntSet
forall a. Monoid a => a
mempty [SCC Vertex]
comps

      neighboursOf :: Vertex -> [Vertex]
neighboursOf Vertex
v = IntGraph
gIntGraph -> Vertex -> [Vertex]
forall i e. Ix i => Array i e -> i -> e
!Vertex
v

      comps :: [SCC Vertex]
comps = IntGraph -> [SCC Vertex]
scc IntGraph
g

      -- To avoid divergence on cyclic input, we build the result
      -- strongly connected component by component, in topological
      -- order. For each SCC, we know that:
      --
      --   * All vertices in the component can reach all other vertices
      --     in the component ("local" reachables)
      --
      --   * Other reachable vertices ("remote" reachables) must come
      --     from earlier components, either via direct neighbourhood, or
      --     transitively from earlier reachability map
      --
      -- This allows us to build the extension of the reachability map
      -- directly, without any self-reference, thereby avoiding a loop.
      add_one_comp :: IM.IntMap IS.IntSet -> SCC Vertex -> IM.IntMap IS.IntSet
      add_one_comp :: IntMap IntSet -> SCC Vertex -> IntMap IntSet
add_one_comp IntMap IntSet
earlier (AcyclicSCC Vertex
v) = Vertex -> IntSet -> IntMap IntSet -> IntMap IntSet
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v IntSet
all_remotes IntMap IntSet
earlier
        where
          earlier_neighbours :: [Vertex]
earlier_neighbours = Vertex -> [Vertex]
neighboursOf Vertex
v
          earlier_further :: [IntSet]
earlier_further = (Vertex -> Maybe IntSet) -> [Vertex] -> [IntSet]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((Vertex -> IntMap IntSet -> Maybe IntSet)
-> IntMap IntSet -> Vertex -> Maybe IntSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Vertex -> IntMap IntSet -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup IntMap IntSet
earlier) [Vertex]
earlier_neighbours
          all_remotes :: IntSet
all_remotes = [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([Vertex] -> IntSet
IS.fromList [Vertex]
earlier_neighbours IntSet -> [IntSet] -> [IntSet]
forall a. a -> [a] -> [a]
: [IntSet]
earlier_further)
      add_one_comp IntMap IntSet
earlier (CyclicSCC [Vertex]
vs) = IntMap IntSet -> IntMap IntSet -> IntMap IntSet
forall a. IntMap a -> IntMap a -> IntMap a
IM.union ([(Vertex, IntSet)] -> IntMap IntSet
forall a. [(Vertex, a)] -> IntMap a
IM.fromList [(Vertex
v, Vertex -> IntSet
local Vertex
v IntSet -> IntSet -> IntSet
`IS.union` IntSet
all_remotes) | Vertex
v <- [Vertex]
vs]) IntMap IntSet
earlier
        where
          all_locals :: IntSet
all_locals = [Vertex] -> IntSet
IS.fromList [Vertex]
vs
          local :: Vertex -> IntSet
local Vertex
v = Vertex -> IntSet -> IntSet
IS.delete Vertex
v IntSet
all_locals
              -- Arguably, for a cyclic SCC we should include each
              -- vertex in its own reachable set. However, this could
              -- lead to a lot of extra pain in client code to avoid
              -- looping when traversing the reachability map.
          all_neighbours :: IntSet
all_neighbours = [Vertex] -> IntSet
IS.fromList ((Vertex -> [Vertex]) -> [Vertex] -> [Vertex]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Vertex -> [Vertex]
neighboursOf [Vertex]
vs)
          earlier_neighbours :: IntSet
earlier_neighbours = IntSet
all_neighbours IntSet -> IntSet -> IntSet
IS.\\ IntSet
all_locals
          earlier_further :: [IntSet]
earlier_further = (Vertex -> Maybe IntSet) -> [Vertex] -> [IntSet]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((Vertex -> IntMap IntSet -> Maybe IntSet)
-> IntMap IntSet -> Vertex -> Maybe IntSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Vertex -> IntMap IntSet -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup IntMap IntSet
earlier) (IntSet -> [Vertex]
IS.toList IntSet
earlier_neighbours)
          all_remotes :: IntSet
all_remotes = [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions (IntSet
earlier_neighbours IntSet -> [IntSet] -> [IntSet]
forall a. a -> [a] -> [a]
: [IntSet]
earlier_further)

--------------------------------------------------------------------------------
-- * Reachability queries
--------------------------------------------------------------------------------

-- | 'allReachable' returns the nodes reachable from the given @root@ node.
--
-- Properties:
--  * The list of nodes /does not/ include the @root@ node!
--  * The list of nodes is deterministically ordered, but according to an
--     internal order determined by the indices attributed to graph nodes.
--  * This function has $O(1)$ complexity.
--
-- If you need a topologically sorted list, consider using the functions exposed from 'GHC.Data.Graph.Directed' on 'Graph' instead.
allReachable :: ReachabilityIndex node -> node {-^ The @root@ node -} -> [node] {-^ All nodes reachable from @root@ -}
allReachable :: forall node. ReachabilityIndex node -> node -> [node]
allReachable (ReachabilityIndex IntMap IntSet
index Vertex -> node
from node -> Maybe Vertex
to) node
root = (Vertex -> node) -> [Vertex] -> [node]
forall a b. (a -> b) -> [a] -> [b]
map Vertex -> node
from [Vertex]
result
  where root_i :: Vertex
root_i = String -> Maybe Vertex -> Vertex
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachableFrom" (node -> Maybe Vertex
to node
root)
        hits :: Maybe IntSet
hits = {-# SCC "allReachable" #-} Vertex -> IntMap IntSet -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
root_i IntMap IntSet
index
        result :: [Vertex]
result = IntSet -> [Vertex]
IS.toList (IntSet -> [Vertex]) -> IntSet -> [Vertex]
forall a b. (a -> b) -> a -> b
$! String -> Maybe IntSet -> IntSet
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachableFrom" Maybe IntSet
hits

-- | 'allReachableMany' returns all nodes reachable from the many given @roots@.
--
-- Properties:
--  * The list of nodes /does not/ include the @roots@ node!
--  * The list of nodes is deterministically ordered, but according to an
--     internal order determined by the indices attributed to graph nodes.
--  * This function has $O(n)$ complexity where $n$ is the number of @roots@.
--
-- If you need a topologically sorted list, consider using the functions
-- exposed from 'GHC.Data.Graph.Directed' on 'Graph' instead ('reachableG').
allReachableMany :: ReachabilityIndex node -> [node] {-^ The @roots@ -} -> [node] {-^ All nodes reachable from all @roots@ -}
allReachableMany :: forall node. ReachabilityIndex node -> [node] -> [node]
allReachableMany (ReachabilityIndex IntMap IntSet
index Vertex -> node
from node -> Maybe Vertex
to) [node]
roots = (Vertex -> node) -> [Vertex] -> [node]
forall a b. (a -> b) -> [a] -> [b]
map Vertex -> node
from (IntSet -> [Vertex]
IS.toList IntSet
hits)
  where roots_i :: [Vertex]
roots_i = [ Vertex
v | Just Vertex
v <- (node -> Maybe Vertex) -> [node] -> [Maybe Vertex]
forall a b. (a -> b) -> [a] -> [b]
map node -> Maybe Vertex
to [node]
roots ]
        hits :: IntSet
hits = {-# SCC "allReachableMany" #-}
               [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([IntSet] -> IntSet) -> [IntSet] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Vertex -> IntSet) -> [Vertex] -> [IntSet]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Maybe IntSet -> IntSet
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachablesG" (Maybe IntSet -> IntSet)
-> (Vertex -> Maybe IntSet) -> Vertex -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vertex -> IntMap IntSet -> Maybe IntSet)
-> IntMap IntSet -> Vertex -> Maybe IntSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Vertex -> IntMap IntSet -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup IntMap IntSet
index) [Vertex]
roots_i

-- | Fast reachability query.
--
-- On graph @g@ with nodes @a@ and @b@, @isReachable g a b@
-- asks whether @b@ can be reached through @g@ starting from @a@.
--
-- Properties:
--  * No self loops, i.e. @isReachable _ a a == False@
--  * This function has $O(1)$ complexity.
isReachable :: ReachabilityIndex node {-^ @g@ -}
            -> node -- ^ @a@
            -> node -- ^ @b@
            -> Bool -- ^ @b@ is reachable from @a@
isReachable :: forall node. ReachabilityIndex node -> node -> node -> Bool
isReachable (ReachabilityIndex IntMap IntSet
index Vertex -> node
_ node -> Maybe Vertex
to) node
a node
b =
    Vertex -> IntSet -> Bool
IS.member Vertex
b_i (IntSet -> Bool) -> IntSet -> Bool
forall a b. (a -> b) -> a -> b
$
    String -> Maybe IntSet -> IntSet
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachable" (Maybe IntSet -> IntSet) -> Maybe IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ Vertex -> IntMap IntSet -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
a_i IntMap IntSet
index
  where a_i :: Vertex
a_i = String -> Maybe Vertex -> Vertex
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachable:node not in graph" (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ node -> Maybe Vertex
to node
a
        b_i :: Vertex
b_i = String -> Maybe Vertex -> Vertex
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachable:node not in graph" (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ node -> Maybe Vertex
to node
b

-- | Fast reachability query with many roots.
--
-- On graph @g@ with many nodes @roots@ and node @b@, @isReachableMany g as b@
-- asks whether @b@ can be reached through @g@ from any of the @roots@.
--
-- Properties:
--  * No self loops, i.e. @isReachableMany _ [a] a == False@
--  * This function is $O(n)$ in the number of roots
isReachableMany :: ReachabilityIndex node -- ^ @g@
                -> [node] -- ^ @roots@
                -> node -- ^ @b@
                -> Bool -- ^ @b@ is reachable from any of the @roots@
isReachableMany :: forall node. ReachabilityIndex node -> [node] -> node -> Bool
isReachableMany (ReachabilityIndex IntMap IntSet
index Vertex -> node
_ node -> Maybe Vertex
to) [node]
roots node
b =
    Vertex -> IntSet -> Bool
IS.member Vertex
b_i (IntSet -> Bool) -> IntSet -> Bool
forall a b. (a -> b) -> a -> b
$
    [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([IntSet] -> IntSet) -> [IntSet] -> IntSet
forall a b. (a -> b) -> a -> b
$
    (Vertex -> IntSet) -> [Vertex] -> [IntSet]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Maybe IntSet -> IntSet
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachablesQuery" (Maybe IntSet -> IntSet)
-> (Vertex -> Maybe IntSet) -> Vertex -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vertex -> IntMap IntSet -> Maybe IntSet)
-> IntMap IntSet -> Vertex -> Maybe IntSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Vertex -> IntMap IntSet -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup IntMap IntSet
index) [Vertex]
roots_i
  where roots_i :: [Vertex]
roots_i = [ Vertex
v | Just Vertex
v <- (node -> Maybe Vertex) -> [node] -> [Maybe Vertex]
forall a b. (a -> b) -> [a] -> [b]
map node -> Maybe Vertex
to [node]
roots ]
        b_i :: Vertex
b_i = String -> Maybe Vertex -> Vertex
forall a. HasDebugCallStack => String -> Maybe a -> a
expectJust String
"reachablesQuery:node not in graph" (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ node -> Maybe Vertex
to node
b