content: -- | 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
  
    -- * Debugging
    , reachabilityIndexMembers
  
    )
    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 {
      index :: IM.IntMap IS.IntSet,
      from_vertex :: Vertex -> node,
      to_vertex :: node -> Maybe Vertex
  }
  
  --
  reachabilityIndexMembers :: ReachabilityIndex node -> [node]
  reachabilityIndexMembers (ReachabilityIndex index from_vert _) = map from_vert (IM.keys index)
  
  --------------------------------------------------------------------------------
  -- * Construction
  --------------------------------------------------------------------------------
  
  -- | Construct a 'ReachabilityIndex' from an acyclic 'Graph'.
  -- If the graph can have cycles, use 'cyclicGraphReachability'
  graphReachability :: Graph node -> ReachabilityIndex node
  graphReachability (Graph g from to) =
    ReachabilityIndex{index = reachableGraph, from_vertex = from, to_vertex = to}
      where
        reachableGraph :: IM.IntMap IS.IntSet
        reachableGraph = IM.fromList [(v, do_one v) | v <- G.vertices g]
  
        do_one v = IS.unions (IS.fromList (g ! v) : mapMaybe (flip IM.lookup reachableGraph) (g ! 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 (Graph g from to) =
    ReachabilityIndex{index = reachableGraphCyclic, from_vertex = from, to_vertex = to}
      where
        reachableGraphCyclic :: IM.IntMap IS.IntSet
        reachableGraphCyclic = foldl' add_one_comp mempty comps
  
        neighboursOf v = g!v
  
        comps = scc 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 earlier (AcyclicSCC v) = IM.insert v all_remotes earlier
          where
            earlier_neighbours = neighboursOf v
            earlier_further = mapMaybe (flip IM.lookup earlier) earlier_neighbours
            all_remotes = IS.unions (IS.fromList earlier_neighbours : earlier_further)
        add_one_comp earlier (CyclicSCC vs) = IM.union (IM.fromList [(v, local v `IS.union` all_remotes) | v <- vs]) earlier
          where
            all_locals = IS.fromList vs
            local v = IS.delete v 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 = IS.fromList (concatMap neighboursOf vs)
            earlier_neighbours = all_neighbours IS.\\ all_locals
            earlier_further = mapMaybe (flip IM.lookup earlier) (IS.toList earlier_neighbours)
            all_remotes = IS.unions (earlier_neighbours : 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.
  --
  -- 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 (ReachabilityIndex index from to) root = map from result
    where root_i = expectJust (to root)
          hits = {-# SCC "allReachable" #-} IM.lookup root_i index
          result = IS.toList $! expectJust 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 (ReachabilityIndex index from to) roots = map from (IS.toList hits)
    where roots_i = [ v | Just v <- map to roots ]
          hits = {-# SCC "allReachableMany" #-}
                 IS.unions $ map (expectJust . flip IM.lookup index) 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@
  isReachable :: ReachabilityIndex node {-^ @g@ -}
              -> node -- ^ @a@
              -> node -- ^ @b@
              -> Bool -- ^ @b@ is reachable from @a@
  isReachable (ReachabilityIndex index _ to) a b =
      IS.member b_i $
      expectJust $ IM.lookup a_i index
    where a_i = expectJust $ to a
          b_i = expectJust $ to 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@.
  --
  -- By partially applying this function to a set of roots, the resulting function can
  -- be applied many times and share the initial work.
  --
  -- Properties:
  --  * No self loops, i.e. @isReachableMany _ [a] a == False@
  isReachableMany :: ReachabilityIndex node -- ^ @g@
                  -> [node] -- ^ @roots@
                  -> (node -> Bool) -- ^ @b@ is reachable from any of the @roots@
  isReachableMany (ReachabilityIndex index _ to) roots =
    let roots_i = [ v | Just v <- map to roots ]
        unions =
            IS.unions $
              map (expectJust . flip IM.lookup index) roots_i
    in \b -> let b_i = expectJust $ to b
             in IS.member b_i unions
  