Merge pull request #887 from meooow25/graph-dfs

Optimize Data.Graph.dfs

Author: treeowl

containers-tests/benchmarks/Graph.hs

@@ -12,19 +12,22 @@ main = do
   evaluate $ rnf randGs
   defaultMain
     [ bgroup "buildG" $ forGs randGs $ \g -> nf (G.buildG (bounds (getG g))) (getEdges g)
-    , bgroup "graphFromEdges" $ forGs randGs $ nf ((\(g, _, _) -> g) . G.graphFromEdges) . getAdjList
+    , bgroup "graphFromEdges" $
+        forGs [randG1, randG2, randG3] $ nf ((\(g, _, _) -> g) . G.graphFromEdges) . getAdjList
     , bgroup "dfs" $ forGs randGs $ nf (flip G.dfs [1]) . getG
     , bgroup "dff" $ forGs randGs $ nf G.dff . getG
     , bgroup "topSort" $ forGs randGs $ nf G.topSort . getG
     , bgroup "scc" $ forGs randGs $ nf G.scc . getG
     , bgroup "bcc" $ forGs [randG1, randG2] $ nf G.bcc . getG
-    , bgroup "stronglyConnCompR" $ forGs randGs $ nf G.stronglyConnCompR . getAdjList
+    , bgroup "stronglyConnCompR" $
+        forGs [randG1, randG2, randG3] $ nf G.stronglyConnCompR . getAdjList
     ]
   where
     randG1 = buildRandG 100 1000 
     randG2 = buildRandG 100 10000
     randG3 = buildRandG 10000 100000
-    randGs = [randG1, randG2, randG3]
+    randG4 = buildRandG 100000 1000000
+    randGs = [randG1, randG2, randG3, randG4]
 
 -- Note: In practice it does not make sense to run topSort or bcc on a random
 -- graph. For topSort the graph should be acyclic and for bcc the graph should

containers/src/Data/Graph.hs

@@ -472,27 +472,24 @@ dff g         = dfs g (vertices g)
 -- | A spanning forest of the part of the graph reachable from the listed
 -- vertices, obtained from a depth-first search of the graph starting at
 -- each of the listed vertices in order.
-dfs          :: Graph -> [Vertex] -> Forest Vertex
-dfs g vs      = prune (bounds g) (map (generate g) vs)
-
-generate     :: Graph -> Vertex -> Tree Vertex
-generate g v  = Node v (map (generate g) (g!v))
-
-prune        :: Bounds -> Forest Vertex -> Forest Vertex
-prune bnds ts = run bnds (chop ts)
-
-chop         :: Forest Vertex -> SetM s (Forest Vertex)
-chop []       = return []
-chop (Node v ts : us)
-              = do
-                visited <- contains v
-                if visited then
-                  chop us
-                 else do
-                  include v
-                  as <- chop ts
-                  bs <- chop us
-                  return (Node v as : bs)
+
+-- This function deviates from King and Launchbury's implementation by
+-- bundling together the functions generate, prune, and chop for efficiency
+-- reasons.
+dfs :: Graph -> [Vertex] -> Forest Vertex
+dfs g vs0 = run (bounds g) $ go vs0
+  where
+    go :: [Vertex] -> SetM s (Forest Vertex)
+    go [] = pure []
+    go (v:vs) = do
+      visited <- contains v
+      if visited
+      then go vs
+      else do
+        include v
+        as <- go (g!v)
+        bs <- go vs
+        pure $ Node v as : bs
 
 -- A monad holding a set of vertices visited so far.
 #if USE_ST_MONAD