Generalise scc and topSort

Author: paf31

.gitignore

@@ -1,3 +1,4 @@
+/output/
 /js/
 /externs/
 /node_modules/

src/Data/Graph.purs.hs

@@ -3,11 +3,14 @@ module Data.Graph (
   Graph(..),
 
   scc,
-  topSort
+  scc',
+  
+  topSort,
+  topSort'
   ) where
 
 import Data.Maybe
-import Data.Array (reverse, concatMap)
+import Data.Array (map, reverse, concatMap)
 import Data.Foldable
 import Data.Traversable
 
@@ -25,7 +28,10 @@ data Graph v = Graph [v] [Edge v]
 type Index = Number
 
 scc :: forall v. (Eq v, Ord v) => Graph v -> [[v]]
-scc (Graph vs es) = runPure (runST (do
+scc = scc' id
+
+scc' :: forall k v. (Eq k, Ord k) => (v -> k) -> Graph v -> [[v]]
+scc' makeKey (Graph vs es) = runPure (runST (do
   index      <- newSTRef 0
   path       <- newSTRef []
   indexMap   <- newSTRef M.empty
@@ -35,11 +41,11 @@ scc (Graph vs es) = runPure (runST (do
   (let 
     indexOf v = do
       m <- readSTRef indexMap
-      return $ M.lookup v m
+      return $ M.lookup (makeKey v) m
 
     lowlinkOf v = do
       m <- readSTRef lowlinkMap
-      return $ M.lookup v m
+      return $ M.lookup (makeKey v) m
 
     go [] = readSTRef components
     go (v : vs) = do
@@ -50,13 +56,13 @@ scc (Graph vs es) = runPure (runST (do
     strongConnect v = do
       i <- readSTRef index
 
-      modifySTRef indexMap   $ M.insert v i
-      modifySTRef lowlinkMap $ M.insert v i
+      modifySTRef indexMap   $ M.insert (makeKey v) i
+      modifySTRef lowlinkMap $ M.insert (makeKey v) i
 
       writeSTRef index $ i + 1
       modifySTRef path $ (:) v
 
-      for es $ \(Edge v' w) -> when (v == v') $ do
+      for es $ \(Edge v' w) -> when (makeKey v == makeKey v') $ do
         wIndex <- indexOf w
         currentPath <- readSTRef path
 
@@ -65,27 +71,27 @@ scc (Graph vs es) = runPure (runST (do
             strongConnect w
             wLowlink <- lowlinkOf w
             for_ wLowlink $ \lowlink ->
-              modifySTRef lowlinkMap $ M.alter (maybeMin lowlink) v
-          _ -> when (w `elem` currentPath) $ do
+              modifySTRef lowlinkMap $ M.alter (maybeMin lowlink) (makeKey v)
+          _ -> when (makeKey w `elem` map makeKey currentPath) $ do
                  wIndex <- indexOf w
                  for_ wIndex $ \index ->
-                   modifySTRef lowlinkMap $ M.alter (maybeMin index) v
+                   modifySTRef lowlinkMap $ M.alter (maybeMin index) (makeKey v)
 
       vIndex <- indexOf v
       vLowlink <- lowlinkOf v        
 
       when (vIndex == vLowlink) $ do
         currentPath <- readSTRef path
-        let newPath = popUntil v currentPath []
+        let newPath = popUntil makeKey v currentPath []
         modifySTRef components $ flip (++) [newPath.component]
         writeSTRef path newPath.path
         return {}
     in go vs)))
 
-popUntil :: forall v. (Eq v) => v -> [v] -> [v] -> { path :: [v], component :: [v] }
-popUntil _ [] popped = { path: [], component: popped } 
-popUntil v (w : path) popped | v == w = { path: path, component: w : popped }
-popUntil v (w : ws) popped = popUntil v ws (w : popped)
+popUntil :: forall k v. (Eq k) => (v -> k) -> v -> [v] -> [v] -> { path :: [v], component :: [v] }
+popUntil _       _ []         popped = { path: [], component: popped } 
+popUntil makeKey v (w : path) popped | makeKey v == makeKey w = { path: path, component: w : popped }
+popUntil makeKey v (w : ws)   popped = popUntil makeKey v ws (w : popped)
 
 maybeMin :: Index -> Maybe Index -> Maybe Index
 maybeMin i Nothing = Just i
@@ -95,4 +101,7 @@ maybeMin i (Just j) = Just $ Math.min i j
 -- Topological sort
 --
 topSort :: forall v. (Eq v, Ord v) => Graph v -> [v]
-topSort = reverse <<< concatMap id <<< scc
+topSort = topSort' id
+
+topSort' :: forall k v. (Eq k, Ord k) => (v -> k) -> Graph v -> [v]
+topSort' makeKey = reverse <<< concatMap id <<< scc' makeKey