-- | Depth-first search and derived operations.
--
-- All of the search variants take a list of 'Vertex' that serves as
-- roots for the search.
--
-- The [x] variants ('xdfsWith' and 'xdffWith') are the most general
-- and are fully configurable in direction and action.  They take a
-- \"direction\" function that tells the search what vertices are
-- next from the current 'Vertex'.  They also take a summarization function
-- to convert a 'Vertex' into some other value.  This could be 'id' or a
-- function to extract a label, if supported by your graph type.
--
-- The [r] variants are reverse searches, while the [u] variants are
-- undirected.
--
-- A depth-first forest is a collection (list) of depth-first trees.  A
-- depth-first tree is an n-ary tree rooted at a vertex that contains
-- the vertices reached in a depth-first search from that root.  The
-- edges in the tree are a subset of the edges in the graph.
module Data.Graph.Haggle.Algorithms.DFS (
  -- * Depth-first Searches
  xdfsWith,
  dfsWith,
  dfs,
  rdfsWith,
  rdfs,
  udfsWith,
  udfs,
  -- * Depth-first Forests
  xdffWith,
  dffWith,
  dff,
  rdffWith,
  rdff,
  udffWith,
  udff,
  -- * Derived Queries
  components,
  noComponents,
  isConnected,
  topsort,
  scc,
  reachable
  ) where

import Control.Monad ( filterM, foldM, liftM )
import Control.Monad.ST
import qualified Data.Foldable as F
import Data.Monoid
import qualified Data.Sequence as Seq
import Data.Tree ( Tree )
import qualified Data.Tree as T

import Prelude

import Data.Graph.Haggle
import Data.Graph.Haggle.Classes ( maxVertexId )
import Data.Graph.Haggle.Internal.BitSet

-- | The most general DFS
xdfsWith :: (Graph g)
         => g
         -> (Vertex -> [Vertex])
         -> (Vertex -> c)
         -> [Vertex]
         -> [c]
xdfsWith :: g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [c]
xdfsWith g
g Vertex -> [Vertex]
nextVerts Vertex -> c
f [Vertex]
roots
  | g -> Bool
forall g. Graph g => g -> Bool
isEmpty g
g Bool -> Bool -> Bool
|| [Vertex] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Vertex]
roots = []
  | Bool
otherwise = (forall s. ST s [c]) -> [c]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [c]) -> [c]) -> (forall s. ST s [c]) -> [c]
forall a b. (a -> b) -> a -> b
$ do
    BitSet s
bs <- Int -> ST s (BitSet s)
forall s. Int -> ST s (BitSet s)
newBitSet (g -> Int
forall g. Graph g => g -> Int
maxVertexId g
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    [c]
res <- ([c] -> Vertex -> ST s [c]) -> [c] -> [Vertex] -> ST s [c]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (BitSet s -> [c] -> Vertex -> ST s [c]
forall s. BitSet s -> [c] -> Vertex -> ST s [c]
go BitSet s
bs) [] [Vertex]
roots
    [c] -> ST s [c]
forall (m :: * -> *) a. Monad m => a -> m a
return ([c] -> ST s [c]) -> [c] -> ST s [c]
forall a b. (a -> b) -> a -> b
$ [c] -> [c]
forall a. [a] -> [a]
reverse [c]
res
  where
    go :: BitSet s -> [c] -> Vertex -> ST s [c]
go BitSet s
bs [c]
acc Vertex
v = do
      Bool
isMarked <- BitSet s -> Int -> ST s Bool
forall s. BitSet s -> Int -> ST s Bool
testBit BitSet s
bs (Vertex -> Int
vertexId Vertex
v)
      case Bool
isMarked of
        Bool
True -> [c] -> ST s [c]
forall (m :: * -> *) a. Monad m => a -> m a
return [c]
acc
        Bool
False -> do
          BitSet s -> Int -> ST s ()
forall s. BitSet s -> Int -> ST s ()
setBit BitSet s
bs (Vertex -> Int
vertexId Vertex
v)
          [Vertex]
nxt <- (Vertex -> ST s Bool) -> [Vertex] -> ST s [Vertex]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (BitSet s -> Vertex -> ST s Bool
forall s. BitSet s -> Vertex -> ST s Bool
notVisited BitSet s
bs) (Vertex -> [Vertex]
nextVerts Vertex
v)
          ([c] -> Vertex -> ST s [c]) -> [c] -> [Vertex] -> ST s [c]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (BitSet s -> [c] -> Vertex -> ST s [c]
go BitSet s
bs) (Vertex -> c
f Vertex
v c -> [c] -> [c]
forall a. a -> [a] -> [a]
: [c]
acc) [Vertex]
nxt

notVisited :: BitSet s -> Vertex -> ST s Bool
notVisited :: BitSet s -> Vertex -> ST s Bool
notVisited BitSet s
bs Vertex
v = (Bool -> Bool) -> ST s Bool -> ST s Bool
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM Bool -> Bool
not (BitSet s -> Int -> ST s Bool
forall s. BitSet s -> Int -> ST s Bool
testBit BitSet s
bs (Vertex -> Int
vertexId Vertex
v))

-- | Forward parameterized DFS
dfsWith :: (Graph g)
        => g
        -> (Vertex -> c)
        -> [Vertex]
        -> [c]
dfsWith :: g -> (Vertex -> c) -> [Vertex] -> [c]
dfsWith g
g = g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [c]
forall g c.
Graph g =>
g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [c]
xdfsWith g
g (g -> Vertex -> [Vertex]
forall g. Graph g => g -> Vertex -> [Vertex]
successors g
g)

-- | Forward DFS
dfs :: (Graph g) => g -> [Vertex] -> [Vertex]
dfs :: g -> [Vertex] -> [Vertex]
dfs g
g = g -> (Vertex -> Vertex) -> [Vertex] -> [Vertex]
forall g c. Graph g => g -> (Vertex -> c) -> [Vertex] -> [c]
dfsWith g
g Vertex -> Vertex
forall a. a -> a
id

-- | Reverse parameterized DFS
rdfsWith :: (Bidirectional g)
         => g
         -> (Vertex -> c)
         -> [Vertex]
         -> [c]
rdfsWith :: g -> (Vertex -> c) -> [Vertex] -> [c]
rdfsWith g
g = g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [c]
forall g c.
Graph g =>
g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [c]
xdfsWith g
g (g -> Vertex -> [Vertex]
forall g. Bidirectional g => g -> Vertex -> [Vertex]
predecessors g
g)

-- | Reverse DFS
rdfs :: (Bidirectional g) => g -> [Vertex] -> [Vertex]
rdfs :: g -> [Vertex] -> [Vertex]
rdfs g
g = g -> (Vertex -> Vertex) -> [Vertex] -> [Vertex]
forall g c.
Bidirectional g =>
g -> (Vertex -> c) -> [Vertex] -> [c]
rdfsWith g
g Vertex -> Vertex
forall a. a -> a
id

-- | Undirected parameterized DFS.  This variant follows both
-- incoming and outgoing edges from each 'Vertex'.
udfsWith :: (Bidirectional g)
         => g
         -> (Vertex -> c)
         -> [Vertex]
         -> [c]
udfsWith :: g -> (Vertex -> c) -> [Vertex] -> [c]
udfsWith g
g = g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [c]
forall g c.
Graph g =>
g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [c]
xdfsWith g
g (g -> Vertex -> [Vertex]
forall g. Bidirectional g => g -> Vertex -> [Vertex]
neighbors g
g)

-- | Undirected DFS
udfs :: (Bidirectional g) => g -> [Vertex] -> [Vertex]
udfs :: g -> [Vertex] -> [Vertex]
udfs g
g = g -> (Vertex -> Vertex) -> [Vertex] -> [Vertex]
forall g c.
Bidirectional g =>
g -> (Vertex -> c) -> [Vertex] -> [c]
udfsWith g
g Vertex -> Vertex
forall a. a -> a
id

-- | The most general depth-first forest.
xdffWith :: (Graph g)
         => g
         -> (Vertex -> [Vertex])
         -> (Vertex -> c)
         -> [Vertex]
         -> [Tree c]
xdffWith :: g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [Tree c]
xdffWith g
g Vertex -> [Vertex]
nextVerts Vertex -> c
f [Vertex]
roots
  | g -> Bool
forall g. Graph g => g -> Bool
isEmpty g
g Bool -> Bool -> Bool
|| [Vertex] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Vertex]
roots = []
  | Bool
otherwise = (forall s. ST s [Tree c]) -> [Tree c]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [Tree c]) -> [Tree c])
-> (forall s. ST s [Tree c]) -> [Tree c]
forall a b. (a -> b) -> a -> b
$ do
    BitSet s
bs <- Int -> ST s (BitSet s)
forall s. Int -> ST s (BitSet s)
newBitSet (g -> Int
forall g. Graph g => g -> Int
maxVertexId g
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    [Tree c]
res <- ([Tree c] -> Vertex -> ST s [Tree c])
-> [Tree c] -> [Vertex] -> ST s [Tree c]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (BitSet s -> [Tree c] -> Vertex -> ST s [Tree c]
forall s. BitSet s -> [Tree c] -> Vertex -> ST s [Tree c]
go BitSet s
bs) [] [Vertex]
roots
    [Tree c] -> ST s [Tree c]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Tree c] -> ST s [Tree c]) -> [Tree c] -> ST s [Tree c]
forall a b. (a -> b) -> a -> b
$ [Tree c] -> [Tree c]
forall a. [a] -> [a]
reverse [Tree c]
res
  where
    go :: BitSet s -> [Tree c] -> Vertex -> ST s [Tree c]
go BitSet s
bs [Tree c]
acc Vertex
v = do
      Bool
isMarked <- BitSet s -> Int -> ST s Bool
forall s. BitSet s -> Int -> ST s Bool
testBit BitSet s
bs (Vertex -> Int
vertexId Vertex
v)
      case Bool
isMarked of
        Bool
True -> [Tree c] -> ST s [Tree c]
forall (m :: * -> *) a. Monad m => a -> m a
return [Tree c]
acc
        Bool
False -> do
          BitSet s -> Int -> ST s ()
forall s. BitSet s -> Int -> ST s ()
setBit BitSet s
bs (Vertex -> Int
vertexId Vertex
v)
          [Vertex]
nxt <- (Vertex -> ST s Bool) -> [Vertex] -> ST s [Vertex]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (BitSet s -> Vertex -> ST s Bool
forall s. BitSet s -> Vertex -> ST s Bool
notVisited BitSet s
bs) (Vertex -> [Vertex]
nextVerts Vertex
v)
          [Tree c]
ts <- ([Tree c] -> Vertex -> ST s [Tree c])
-> [Tree c] -> [Vertex] -> ST s [Tree c]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (BitSet s -> [Tree c] -> Vertex -> ST s [Tree c]
go BitSet s
bs) [] [Vertex]
nxt
          [Tree c] -> ST s [Tree c]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Tree c] -> ST s [Tree c]) -> [Tree c] -> ST s [Tree c]
forall a b. (a -> b) -> a -> b
$ c -> [Tree c] -> Tree c
forall a. a -> Forest a -> Tree a
T.Node (Vertex -> c
f Vertex
v) ([Tree c] -> [Tree c]
forall a. [a] -> [a]
reverse [Tree c]
ts) Tree c -> [Tree c] -> [Tree c]
forall a. a -> [a] -> [a]
: [Tree c]
acc

dffWith :: (Graph g)
        => g
        -> (Vertex -> c)
        -> [Vertex]
        -> [Tree c]
dffWith :: g -> (Vertex -> c) -> [Vertex] -> [Tree c]
dffWith g
g = g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [Tree c]
forall g c.
Graph g =>
g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [Tree c]
xdffWith g
g (g -> Vertex -> [Vertex]
forall g. Graph g => g -> Vertex -> [Vertex]
successors g
g)

dff :: (Graph g) => g -> [Vertex] -> [Tree Vertex]
dff :: g -> [Vertex] -> [Tree Vertex]
dff g
g = g -> (Vertex -> Vertex) -> [Vertex] -> [Tree Vertex]
forall g c. Graph g => g -> (Vertex -> c) -> [Vertex] -> [Tree c]
dffWith g
g Vertex -> Vertex
forall a. a -> a
id

rdffWith :: (Bidirectional g) => g -> (Vertex -> c) -> [Vertex] -> [Tree c]
rdffWith :: g -> (Vertex -> c) -> [Vertex] -> [Tree c]
rdffWith g
g = g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [Tree c]
forall g c.
Graph g =>
g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [Tree c]
xdffWith g
g (g -> Vertex -> [Vertex]
forall g. Bidirectional g => g -> Vertex -> [Vertex]
predecessors g
g)

rdff :: (Bidirectional g) => g -> [Vertex] -> [Tree Vertex]
rdff :: g -> [Vertex] -> [Tree Vertex]
rdff g
g = g -> (Vertex -> Vertex) -> [Vertex] -> [Tree Vertex]
forall g c.
Bidirectional g =>
g -> (Vertex -> c) -> [Vertex] -> [Tree c]
rdffWith g
g Vertex -> Vertex
forall a. a -> a
id

udffWith :: (Bidirectional g) => g -> (Vertex -> c) -> [Vertex] -> [Tree c]
udffWith :: g -> (Vertex -> c) -> [Vertex] -> [Tree c]
udffWith g
g = g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [Tree c]
forall g c.
Graph g =>
g -> (Vertex -> [Vertex]) -> (Vertex -> c) -> [Vertex] -> [Tree c]
xdffWith g
g (g -> Vertex -> [Vertex]
forall g. Bidirectional g => g -> Vertex -> [Vertex]
neighbors g
g)

udff :: (Bidirectional g) => g -> [Vertex] -> [Tree Vertex]
udff :: g -> [Vertex] -> [Tree Vertex]
udff g
g = g -> (Vertex -> Vertex) -> [Vertex] -> [Tree Vertex]
forall g c.
Bidirectional g =>
g -> (Vertex -> c) -> [Vertex] -> [Tree c]
udffWith g
g Vertex -> Vertex
forall a. a -> a
id

-- Derived

-- | Return a list of each connected component in the graph
components :: (Bidirectional g) => g -> [[Vertex]]
components :: g -> [[Vertex]]
components g
g = (Tree Vertex -> [Vertex]) -> [Tree Vertex] -> [[Vertex]]
forall a b. (a -> b) -> [a] -> [b]
map Tree Vertex -> [Vertex]
forall a. Tree a -> [a]
preorder ([Tree Vertex] -> [[Vertex]]) -> [Tree Vertex] -> [[Vertex]]
forall a b. (a -> b) -> a -> b
$ g -> [Vertex] -> [Tree Vertex]
forall g. Bidirectional g => g -> [Vertex] -> [Tree Vertex]
udff g
g (g -> [Vertex]
forall g. Graph g => g -> [Vertex]
vertices g
g)

-- | The number of components in the graph
noComponents :: (Bidirectional g) => g -> Int
noComponents :: g -> Int
noComponents = [[Vertex]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[Vertex]] -> Int) -> (g -> [[Vertex]]) -> g -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> [[Vertex]]
forall g. Bidirectional g => g -> [[Vertex]]
components

-- | True if there is only a single component in the graph.
isConnected :: (Bidirectional g) => g -> Bool
isConnected :: g -> Bool
isConnected = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
1) (Int -> Bool) -> (g -> Int) -> g -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> Int
forall g. Bidirectional g => g -> Int
noComponents

-- | Topologically sort the graph; the input must be a DAG.
topsort :: (Graph g) => g -> [Vertex]
topsort :: g -> [Vertex]
topsort g
g = [Vertex] -> [Vertex]
forall a. [a] -> [a]
reverse ([Vertex] -> [Vertex]) -> [Vertex] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ Seq Vertex -> [Vertex]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (Seq Vertex -> [Vertex]) -> Seq Vertex -> [Vertex]
forall a b. (a -> b) -> a -> b
$ [Tree Vertex] -> Seq Vertex
forall a. [Tree a] -> Seq a
postflattenF ([Tree Vertex] -> Seq Vertex) -> [Tree Vertex] -> Seq Vertex
forall a b. (a -> b) -> a -> b
$ g -> [Vertex] -> [Tree Vertex]
forall g. Graph g => g -> [Vertex] -> [Tree Vertex]
dff g
g (g -> [Vertex]
forall g. Graph g => g -> [Vertex]
vertices g
g)

-- | Return a list of each /strongly-connected component/ in the graph.
-- In a strongly-connected component, every vertex is reachable from every
-- other vertex.
scc :: (Bidirectional g) => g -> [[Vertex]]
scc :: g -> [[Vertex]]
scc g
g = (Tree Vertex -> [Vertex]) -> [Tree Vertex] -> [[Vertex]]
forall a b. (a -> b) -> [a] -> [b]
map Tree Vertex -> [Vertex]
forall a. Tree a -> [a]
preorder (g -> [Vertex] -> [Tree Vertex]
forall g. Bidirectional g => g -> [Vertex] -> [Tree Vertex]
rdff g
g (g -> [Vertex]
forall g. Graph g => g -> [Vertex]
topsort g
g))

-- | Compute the set of vertices reachable from a root 'Vertex'.
reachable :: (Graph g) => Vertex -> g -> [Vertex]
reachable :: Vertex -> g -> [Vertex]
reachable Vertex
v g
g = [Tree Vertex] -> [Vertex]
forall a. [Tree a] -> [a]
preorderF (g -> [Vertex] -> [Tree Vertex]
forall g. Graph g => g -> [Vertex] -> [Tree Vertex]
dff g
g [Vertex
v])

-- Helpers

neighbors :: (Bidirectional g) => g -> Vertex -> [Vertex]
neighbors :: g -> Vertex -> [Vertex]
neighbors g
g Vertex
v = g -> Vertex -> [Vertex]
forall g. Graph g => g -> Vertex -> [Vertex]
successors g
g Vertex
v [Vertex] -> [Vertex] -> [Vertex]
forall a. [a] -> [a] -> [a]
++ g -> Vertex -> [Vertex]
forall g. Bidirectional g => g -> Vertex -> [Vertex]
predecessors g
g Vertex
v

preorder :: Tree a -> [a]
preorder :: Tree a -> [a]
preorder = Tree a -> [a]
forall a. Tree a -> [a]
T.flatten

preorderF :: [Tree a] -> [a]
preorderF :: [Tree a] -> [a]
preorderF = (Tree a -> [a]) -> [Tree a] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [a]
forall a. Tree a -> [a]
preorder

postflatten :: Tree a -> Seq.Seq a
postflatten :: Tree a -> Seq a
postflatten (T.Node a
v Forest a
ts) = Forest a -> Seq a
forall a. [Tree a] -> Seq a
postflattenF Forest a
ts Seq a -> Seq a -> Seq a
forall a. Semigroup a => a -> a -> a
<> a -> Seq a
forall a. a -> Seq a
Seq.singleton a
v

postflattenF :: [Tree a] -> Seq.Seq a
postflattenF :: [Tree a] -> Seq a
postflattenF = (Tree a -> Seq a) -> [Tree a] -> Seq a
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap Tree a -> Seq a
forall a. Tree a -> Seq a
postflatten