{-# LANGUAGE FlexibleContexts #-} -- | A wrapper around the types and functions from "Data.Graph" to make programming with them less painful. Also -- implements some extra useful goodies such as 'successors' and 'sccGraph', and improves the documentation of -- the behaviour of some functions. -- -- As it wraps "Data.Graph", this module only supports directed graphs with unlabelled edges. -- -- Incorporates code from the 'containers' package which is (c) The University of Glasgow 2002 and based -- on code described in: -- -- /Lazy Depth-First Search and Linear Graph Algorithms in Haskell/, -- by David King and John Launchbury module Data.Graph.Wrapper ( Edge, Graph, vertex, fromListSimple, fromList, fromListLenient, fromListBy, fromVerticesEdges, toList, vertices, edges, successors, outdegree, indegree, transpose, reachableVertices, hasPath, topologicalSort, depthNumbering, SCC(..), stronglyConnectedComponents, sccGraph, traverseWithKey ) where import Data.Graph.Wrapper.Internal import Control.Arrow (second) import Control.Monad import Control.Monad.ST import Data.Array import Data.Array.ST import qualified Data.Graph as G import qualified Data.IntSet as IS import Data.List (sortBy, mapAccumL) import Data.Maybe (fromMaybe, fromJust, mapMaybe) import qualified Data.Map as M import Data.Ord import qualified Data.Set as S import qualified Data.Foldable as Foldable import qualified Data.Traversable as Traversable fst3 :: (a, b, c) -> a fst3 (a, _, _) = a snd3 :: (a, b, c) -> b snd3 (_, b, _) = b thd3 :: (a, b, c) -> c thd3 (_, _, c) = c -- amapWithKey :: Ix i => (i -> v -> v') -> Array i v -> Array i v' -- -- More efficient, but not portable (uses GHC.Arr exports): -- --amapWithKey f arr = unsafeArray' (bounds arr) (numElements arr) [(i, f i (unsafeAt arr i)) | i <- [0 .. n - 1]] -- amapWithKey f arr = array (bounds arr) [(i, f i v) | (i, v) <- assocs arr] amapWithKeyM :: (Monad m, Ix i) => (i -> v -> m v') -> Array i v -> m (Array i v') amapWithKeyM f arr = liftM (array (bounds arr)) $ mapM (\(i, v) -> liftM (\v' -> (i, v')) $ f i v) (assocs arr) -- | Construct a 'Graph' where the vertex data double up as the indices. -- -- Unlike 'Data.Graph.graphFromEdges', vertex data that is listed as edges that are not actually themselves -- present in the input list are reported as an error. fromListSimple :: Ord v => [(v, [v])] -> Graph v v fromListSimple = fromListBy id -- | Construct a 'Graph' that contains the given vertex data, linked up according to the supplied key extraction -- function and edge list. -- -- Unlike 'Data.Graph.graphFromEdges', indexes in the edge list that do not correspond to the index of some item in the -- input list are reported as an error. fromListBy :: Ord i => (v -> i) -> [(v, [i])] -> Graph i v fromListBy f vertices = fromList [(f v, v, is) | (v, is) <- vertices] -- | Construct a 'Graph' directly from a list of vertices (and vertex data). -- -- If either end of an 'Edge' does not correspond to a supplied vertex, an error will be raised. fromVerticesEdges :: Ord i => [(i, v)] -> [Edge i] -> Graph i v fromVerticesEdges vertices edges | M.null final_edges_map = fromList done_vertices | otherwise = error "fromVerticesEdges: some edges originated from non-existant vertices" where (final_edges_map, done_vertices) = mapAccumL accum (M.fromListWith (++) (map (second return) edges)) vertices accum edges_map (i, v) = case M.updateLookupWithKey (\_ _ -> Nothing) i edges_map of (mb_is, edges_map) -> (edges_map, (i, v, fromMaybe [] mb_is)) -- | Construct a 'Graph' that contains the given vertex data, linked up according to the supplied index and edge list. -- -- Unlike 'Data.Graph.graphFromEdges', indexes in the edge list that do not correspond to the index of some item in the -- input list are reported as an error. fromList :: Ord i => [(i, v, [i])] -> Graph i v fromList = fromList' False -- | Construct a 'Graph' that contains the given vertex data, linked up according to the supplied index and edge list. -- -- Like 'Data.Graph.graphFromEdges', indexes in the edge list that do not correspond to the index of some item in the -- input list are silently ignored. fromListLenient :: Ord i => [(i, v, [i])] -> Graph i v fromListLenient = fromList' True {-# INLINE fromList' #-} fromList' :: Ord i => Bool -> [(i, v, [i])] -> Graph i v fromList' lenient vertices = G graph key_map vertex_map where max_v = length vertices - 1 bounds0 = (0, max_v) :: (G.Vertex, G.Vertex) sorted_vertices = sortBy (comparing fst3) vertices index_vertex = if lenient then mapMaybe (indexGVertex'_maybe key_map) else map (indexGVertex' key_map) graph = array bounds0 $ [0..] `zip` map (index_vertex . thd3) sorted_vertices key_map = array bounds0 $ [0..] `zip` map fst3 sorted_vertices vertex_map = array bounds0 $ [0..] `zip` map snd3 sorted_vertices -- | Morally, the inverse of 'fromList'. The order of the elements in the output list is unspecified, as is the order of the edges -- in each node's adjacency list. For this reason, @toList . fromList@ is not necessarily the identity function. toList :: Ord i => Graph i v -> [(i, v, [i])] toList g = [(indexGVertexArray g ! m, gVertexVertexArray g ! m, map (indexGVertexArray g !) ns) | (m, ns) <- assocs (graph g)] -- | Find the vertices we can reach from a vertex with the given indentity successors :: Ord i => Graph i v -> i -> [i] successors g i = map (gVertexIndex g) (graph g ! indexGVertex g i) -- | Number of edges going out of the vertex. -- -- It is worth sharing a partial application of 'outdegree' to the 'Graph' argument if you intend to query -- for the outdegrees of a number of vertices. outdegree :: Ord i => Graph i v -> i -> Int outdegree g = \i -> outdegrees ! indexGVertex g i where outdegrees = G.outdegree (graph g) -- | Number of edges going in to the vertex. -- -- It is worth sharing a partial application of 'indegree' to the 'Graph' argument if you intend to query -- for the indegrees of a number of vertices. indegree :: Ord i => Graph i v -> i -> Int indegree g = \i -> indegrees ! indexGVertex g i where indegrees = G.indegree (graph g) -- | The graph formed by flipping all the edges, so edges from i to j now go from j to i transpose :: Graph i v -> Graph i v transpose g = g { graph = G.transposeG (graph g) } -- | Topological sort of of the graph (<http://en.wikipedia.org/wiki/Topological_sort>). If the graph is acyclic, -- vertices will only appear in the list once all of those vertices with arrows to them have already appeared. -- -- Vertex /i/ precedes /j/ in the output whenever /j/ is reachable from /i/ but not vice versa. topologicalSort :: Graph i v -> [i] topologicalSort g = map (gVertexIndex g) $ G.topSort (graph g) -- | List all of the vertices reachable from the given starting point reachableVertices :: Ord i => Graph i v -> i -> [i] reachableVertices g i = if i `elem` vertices g then map (gVertexIndex g) $ G.reachable (graph g) $ indexGVertex g i else [] -- | Is the second vertex reachable by following edges from the first vertex? -- -- It is worth sharing a partial application of 'hasPath' to the first vertex if you are testing for several -- vertices being reachable from it. hasPath :: Ord i => Graph i v -> i -> i -> Bool hasPath g i1 = (`elem` reachableVertices g i1) -- | Number the vertices in the graph by how far away they are from the given roots. The roots themselves have depth 0, -- and every subsequent link we traverse adds 1 to the depth. If a vertex is not reachable it will have a depth of 'Nothing'. depthNumbering :: Ord i => Graph i v -> [i] -> Graph i (v, Maybe Int) depthNumbering g is = runST $ do -- This array records the minimum known depth for the node at the moment depth_array <- newArray (bounds (graph g)) Nothing :: ST s (STArray s G.Vertex (Maybe Int)) let -- Lets us adjust the known depth given a new observation atDepth gv depth = do mb_old_depth <- readArray depth_array gv let depth' = maybe depth (`min` depth) mb_old_depth depth' `seq` writeArray depth_array gv (Just depth') -- Do an depth-first search on the graph (checking for cycles to prevent non-termination), -- recording the depth at which any node was seen in that array. let gos seen depth gvs = mapM_ (go seen depth) gvs go seen depth gv | depth `seq` False = error "depthNumbering: unreachable" | gv `IS.member` seen = return () | otherwise = do gv `atDepth` depth gos (IS.insert gv seen) (depth + 1) (graph g ! gv) gos IS.empty 0 (map (indexGVertex g) is) -- let go _ _ [] = return () -- go seen depth gvs = do -- let go_one (seen, next_gvs) gv -- | gv `IS.member` seen = return (seen, next_gvs) -- | otherwise = do gv `atDepth` depth -- return (IS.insert gv seen, next_gvs ++ (graph g ! gv)) -- (seen, next_gvs) <- foldM go_one (seen, []) gvs -- go seen (depth + 1) next_gvs -- -- go IS.empty 0 (map (indexGVertex g) is) gvva <- amapWithKeyM (\gv v -> liftM (\mb_depth -> (v, mb_depth)) $ readArray depth_array gv) (gVertexVertexArray g) return $ g { gVertexVertexArray = gvva } data SCC i = AcyclicSCC i | CyclicSCC [i] deriving (Show) instance Functor SCC where fmap f (AcyclicSCC v) = AcyclicSCC (f v) fmap f (CyclicSCC vs) = CyclicSCC (map f vs) instance Foldable.Foldable SCC where foldMap f (AcyclicSCC v) = f v foldMap f (CyclicSCC vs) = Foldable.foldMap f vs instance Traversable.Traversable SCC where traverse f (AcyclicSCC v) = fmap AcyclicSCC (f v) traverse f (CyclicSCC vs) = fmap CyclicSCC (Traversable.traverse f vs) -- | Strongly connected components (<http://en.wikipedia.org/wiki/Strongly_connected_component>). -- -- The SCCs are listed in a *reverse topological order*. That is to say, any edges *to* a node in the SCC -- originate either *from*: -- -- 1) Within the SCC itself (in the case of a 'CyclicSCC' only) -- 2) Or from a node in a SCC later on in the list -- -- Vertex /i/ strictly precedes /j/ in the output whenever /i/ is reachable from /j/ but not vice versa. -- Vertex /i/ occurs in the same SCC as /j/ whenever both /i/ is reachable from /j/ and /j/ is reachable from /i/. stronglyConnectedComponents :: Graph i v -> [SCC i] stronglyConnectedComponents g = map decode forest where forest = G.scc (graph g) decode (G.Node v []) | mentions_itself v = CyclicSCC [gVertexIndex g v] | otherwise = AcyclicSCC (gVertexIndex g v) decode other = CyclicSCC (dec other []) where dec (G.Node v ts) vs = gVertexIndex g v : foldr dec vs ts mentions_itself v = v `elem` (graph g ! v) -- | The graph formed by the strongly connected components of the input graph. Each node in the resulting -- graph is indexed by the set of vertex indices from the input graph that it contains. sccGraph :: Ord i => Graph i v -> Graph (S.Set i) (M.Map i v) sccGraph g = fromList nodes' where -- As we consume the SCCs, we accumulate a Map i (S.Set i) that tells us which SCC any given index belongs to. -- When we do a lookup, it is sufficient to look in the map accumulated so far because nodes that are successors -- of a SCC must occur to the *left* of it in the list. (_final_i2scc_i, nodes') = mapAccumL go M.empty (stronglyConnectedComponents g) --go :: M.Map i (S.Set i) -> SCC i -> (M.Map i (S.Set i), (S.Set i, M.Map i v, [S.Set i])) go i2scc_i scc = (i2scc_i', (scc_i, Foldable.foldMap (\i -> M.singleton i (vertex g i)) scc, Foldable.foldMap (\i -> map (fromJust . (`M.lookup` i2scc_i')) (successors g i)) scc)) where -- The mechanism by which we index the new graph -- the set of indexes of its components scc_i = Foldable.foldMap S.singleton scc i2scc_i' = i2scc_i `M.union` Foldable.foldMap (\i -> M.singleton i scc_i) scc