{-# OPTIONS_GHC -Wall #-}

module Data.Graph.Good
  ( Graph
  , graphFromEdges
  , vertices
  , edges
  , outdegree
  , indegree
  , transposeG
  , dfs
  , dff
  , topSort
  , reverseTopSort
  , components
  , scc
  , bcc
  , reachable
  , path
  ) where

import           Control.Applicative (empty)
import           Control.Arrow ((***))
import           Control.Monad ((<=<))
import           Data.Array (Ix, Array)
import qualified Data.Array as A
import qualified Data.Graph as G
import           Data.Maybe (mapMaybe, fromMaybe)


data Graph v = Graph
  { Graph v -> Graph
g_graph :: G.Graph
  , Graph v -> Vertex -> v
g_from_vert :: G.Vertex -> v
  , Graph v -> v -> Maybe Vertex
g_to_vert :: v -> Maybe G.Vertex
  }


graphFromEdges :: Ord v => [(v, [v])] -> Graph v
graphFromEdges :: [(v, [v])] -> Graph v
graphFromEdges [(v, [v])]
vs =
  let (Graph
g, Vertex -> (v, v, [v])
v_func, v -> Maybe Vertex
l) = [(v, v, [v])] -> (Graph, Vertex -> (v, v, [v]), v -> Maybe Vertex)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)
G.graphFromEdges ([(v, v, [v])]
 -> (Graph, Vertex -> (v, v, [v]), v -> Maybe Vertex))
-> [(v, v, [v])]
-> (Graph, Vertex -> (v, v, [v]), v -> Maybe Vertex)
forall a b. (a -> b) -> a -> b
$ ((v, [v]) -> (v, v, [v])) -> [(v, [v])] -> [(v, v, [v])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(v
v, [v]
es) -> (v
v, v
v, [v]
es)) [(v, [v])]
vs
   in Graph -> (Vertex -> v) -> (v -> Maybe Vertex) -> Graph v
forall v. Graph -> (Vertex -> v) -> (v -> Maybe Vertex) -> Graph v
Graph Graph
g (\Vertex
vert -> let (v
v, v
_, [v]
_) = Vertex -> (v, v, [v])
v_func Vertex
vert in v
v) v -> Maybe Vertex
l


vertices :: Graph v -> [v]
vertices :: Graph v -> [v]
vertices Graph v
g = Graph v -> [Vertex] -> [v]
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g ([Vertex] -> [v]) -> [Vertex] -> [v]
forall a b. (a -> b) -> a -> b
$ (Graph -> [Vertex]) -> Graph v -> [Vertex]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Vertex]
G.vertices Graph v
g


edges :: Graph v -> [(v, v)]
edges :: Graph v -> [(v, v)]
edges Graph v
g = (Edge -> (v, v)) -> [Edge] -> [(v, v)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Graph v -> Vertex -> v
forall v. Graph v -> Vertex -> v
g_from_vert Graph v
g (Vertex -> v) -> (Vertex -> v) -> Edge -> (v, v)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Graph v -> Vertex -> v
forall v. Graph v -> Vertex -> v
g_from_vert Graph v
g) ([Edge] -> [(v, v)]) -> [Edge] -> [(v, v)]
forall a b. (a -> b) -> a -> b
$ (Graph -> [Edge]) -> Graph v -> [Edge]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Edge]
G.edges Graph v
g


overGraph :: (G.Graph -> r) -> Graph v -> r
overGraph :: (Graph -> r) -> Graph v -> r
overGraph Graph -> r
f = Graph -> r
f (Graph -> r) -> (Graph v -> Graph) -> Graph v -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph v -> Graph
forall v. Graph v -> Graph
g_graph


lookupArr :: Ix k => Array k v -> k -> Maybe v
lookupArr :: Array k v -> k -> Maybe v
lookupArr Array k v
arr k
ix =
  let (k
lo, k
hi) = Array k v -> (k, k)
forall i e. Array i e -> (i, i)
A.bounds Array k v
arr
   in case (k
lo k -> k -> Bool
forall a. Ord a => a -> a -> Bool
<= k
ix Bool -> Bool -> Bool
&& k
ix k -> k -> Bool
forall a. Ord a => a -> a -> Bool
<= k
hi) of
        Bool
True -> v -> Maybe v
forall a. a -> Maybe a
Just (v -> Maybe v) -> v -> Maybe v
forall a b. (a -> b) -> a -> b
$ Array k v
arr Array k v -> k -> v
forall i e. Ix i => Array i e -> i -> e
A.! k
ix
        Bool
False -> Maybe v
forall a. Maybe a
Nothing


outdegree :: Graph v -> v -> Maybe Int
outdegree :: Graph v -> v -> Maybe Vertex
outdegree Graph v
g = Array Vertex Vertex -> Vertex -> Maybe Vertex
forall k v. Ix k => Array k v -> k -> Maybe v
lookupArr Array Vertex Vertex
arr (Vertex -> Maybe Vertex)
-> (v -> Maybe Vertex) -> v -> Maybe Vertex
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Graph v -> v -> Maybe Vertex
forall v. Graph v -> v -> Maybe Vertex
g_to_vert Graph v
g
  where
    arr :: Array Vertex Vertex
arr = (Graph -> Array Vertex Vertex) -> Graph v -> Array Vertex Vertex
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> Array Vertex Vertex
G.outdegree Graph v
g


indegree :: Graph v -> v -> Maybe Int
indegree :: Graph v -> v -> Maybe Vertex
indegree Graph v
g = Array Vertex Vertex -> Vertex -> Maybe Vertex
forall k v. Ix k => Array k v -> k -> Maybe v
lookupArr Array Vertex Vertex
arr (Vertex -> Maybe Vertex)
-> (v -> Maybe Vertex) -> v -> Maybe Vertex
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Graph v -> v -> Maybe Vertex
forall v. Graph v -> v -> Maybe Vertex
g_to_vert Graph v
g
  where
    arr :: Array Vertex Vertex
arr = (Graph -> Array Vertex Vertex) -> Graph v -> Array Vertex Vertex
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> Array Vertex Vertex
G.indegree Graph v
g


transposeG :: Graph v -> Graph v
transposeG :: Graph v -> Graph v
transposeG Graph v
g = Graph v
g { g_graph :: Graph
g_graph = (Graph -> Graph) -> Graph v -> Graph
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> Graph
G.transposeG Graph v
g }


fromVertices :: Functor f => Graph v -> f G.Vertex -> f v
fromVertices :: Graph v -> f Vertex -> f v
fromVertices = (Vertex -> v) -> f Vertex -> f v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Vertex -> v) -> f Vertex -> f v)
-> (Graph v -> Vertex -> v) -> Graph v -> f Vertex -> f v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph v -> Vertex -> v
forall v. Graph v -> Vertex -> v
g_from_vert


dfs :: Graph v -> [v] -> G.Forest v
dfs :: Graph v -> [v] -> Forest v
dfs Graph v
g [v]
vs =
  let verts :: [Vertex]
verts = (v -> Maybe Vertex) -> [v] -> [Vertex]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Graph v -> v -> Maybe Vertex
forall v. Graph v -> v -> Maybe Vertex
g_to_vert Graph v
g) [v]
vs
   in (Tree Vertex -> Tree v) -> [Tree Vertex] -> Forest v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Graph v -> Tree Vertex -> Tree v
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g) ([Tree Vertex] -> Forest v) -> [Tree Vertex] -> Forest v
forall a b. (a -> b) -> a -> b
$ (Graph -> [Vertex] -> [Tree Vertex])
-> Graph v -> [Vertex] -> [Tree Vertex]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Vertex] -> [Tree Vertex]
G.dfs Graph v
g [Vertex]
verts


dff :: Graph v -> G.Forest v
dff :: Graph v -> Forest v
dff Graph v
g = (Tree Vertex -> Tree v) -> [Tree Vertex] -> Forest v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Graph v -> Tree Vertex -> Tree v
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g) ([Tree Vertex] -> Forest v) -> [Tree Vertex] -> Forest v
forall a b. (a -> b) -> a -> b
$ (Graph -> [Tree Vertex]) -> Graph v -> [Tree Vertex]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Tree Vertex]
G.dff Graph v
g


topSort :: Graph v -> [v]
topSort :: Graph v -> [v]
topSort Graph v
g = Graph v -> [Vertex] -> [v]
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g ([Vertex] -> [v]) -> [Vertex] -> [v]
forall a b. (a -> b) -> a -> b
$ (Graph -> [Vertex]) -> Graph v -> [Vertex]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Vertex]
G.topSort Graph v
g


reverseTopSort :: Graph v -> [v]
reverseTopSort :: Graph v -> [v]
reverseTopSort = [v] -> [v]
forall a. [a] -> [a]
reverse ([v] -> [v]) -> (Graph v -> [v]) -> Graph v -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph v -> [v]
forall v. Graph v -> [v]
topSort


components :: Graph v -> G.Forest v
components :: Graph v -> Forest v
components Graph v
g = (Tree Vertex -> Tree v) -> [Tree Vertex] -> Forest v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Graph v -> Tree Vertex -> Tree v
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g) ([Tree Vertex] -> Forest v) -> [Tree Vertex] -> Forest v
forall a b. (a -> b) -> a -> b
$ (Graph -> [Tree Vertex]) -> Graph v -> [Tree Vertex]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Tree Vertex]
G.components Graph v
g


scc :: Graph v -> G.Forest v
scc :: Graph v -> Forest v
scc Graph v
g = (Tree Vertex -> Tree v) -> [Tree Vertex] -> Forest v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Graph v -> Tree Vertex -> Tree v
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g) ([Tree Vertex] -> Forest v) -> [Tree Vertex] -> Forest v
forall a b. (a -> b) -> a -> b
$ (Graph -> [Tree Vertex]) -> Graph v -> [Tree Vertex]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Tree Vertex]
G.scc Graph v
g


bcc :: Graph v -> G.Forest [v]
bcc :: Graph v -> Forest [v]
bcc Graph v
g = (Tree [Vertex] -> Tree [v]) -> [Tree [Vertex]] -> Forest [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Vertex] -> [v]) -> Tree [Vertex] -> Tree [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Vertex] -> [v]) -> Tree [Vertex] -> Tree [v])
-> ([Vertex] -> [v]) -> Tree [Vertex] -> Tree [v]
forall a b. (a -> b) -> a -> b
$ Graph v -> [Vertex] -> [v]
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g) ([Tree [Vertex]] -> Forest [v]) -> [Tree [Vertex]] -> Forest [v]
forall a b. (a -> b) -> a -> b
$ (Graph -> [Tree [Vertex]]) -> Graph v -> [Tree [Vertex]]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> [Tree [Vertex]]
G.bcc Graph v
g


reachable :: Graph v -> v -> [v]
reachable :: Graph v -> v -> [v]
reachable Graph v
g v
v = case Graph v -> v -> Maybe Vertex
forall v. Graph v -> v -> Maybe Vertex
g_to_vert Graph v
g v
v of
  Maybe Vertex
Nothing -> [v]
forall (f :: * -> *) a. Alternative f => f a
empty
  Just Vertex
vert -> Graph v -> [Vertex] -> [v]
forall (f :: * -> *) v. Functor f => Graph v -> f Vertex -> f v
fromVertices Graph v
g ([Vertex] -> [v]) -> [Vertex] -> [v]
forall a b. (a -> b) -> a -> b
$ (Graph -> Vertex -> [Vertex]) -> Graph v -> Vertex -> [Vertex]
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> Vertex -> [Vertex]
G.reachable Graph v
g Vertex
vert


path :: Graph v -> v -> v -> Bool
path :: Graph v -> v -> v -> Bool
path Graph v
g v
v1 v
v2 = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
False (Maybe Bool -> Bool) -> Maybe Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
  Vertex
vert1 <- Graph v -> v -> Maybe Vertex
forall v. Graph v -> v -> Maybe Vertex
g_to_vert Graph v
g v
v1
  Vertex
vert2 <- Graph v -> v -> Maybe Vertex
forall v. Graph v -> v -> Maybe Vertex
g_to_vert Graph v
g v
v2
  Bool -> Maybe Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> Maybe Bool) -> Bool -> Maybe Bool
forall a b. (a -> b) -> a -> b
$ (Graph -> Vertex -> Vertex -> Bool)
-> Graph v -> Vertex -> Vertex -> Bool
forall r v. (Graph -> r) -> Graph v -> r
overGraph Graph -> Vertex -> Vertex -> Bool
G.path Graph v
g Vertex
vert1 Vertex
vert2