Alternative.hs

module Test.Data.Graph.Alternative where

import qualified Data.Graph.Comfort as Graph
import Data.Graph.Comfort (Graph)

import qualified Control.Monad.Trans.State as MS
import Control.Monad (when)
import Control.Applicative (liftA2)

import qualified Data.NonEmpty.Set as NonEmptySet
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Tree as Tree
import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import Data.Map (Map)
import Data.Set (Set)
import Data.Tree (Tree, Forest)
import Data.Maybe (catMaybes)
import Data.Tuple.HT (mapFst, mapSnd)



depthFirstSearch ::
   (Graph.Edge edge, Ord node) =>
   Graph edge node edgeLabel nodeLabel -> Forest node
depthFirstSearch gr =
   let go = do
         unvisited <- MS.get
         case Set.toAscList unvisited of
            n:_ -> liftA2 (:) (depthFirstSearchFrom gr n) go
            [] -> pure []
   in MS.evalState go $ Graph.nodeSet gr

depthFirstSearchFrom ::
   (Graph.Edge edge, Ord node) =>
   Graph edge node edgeLabel nodeLabel ->
   node -> MS.State (Set node) (Tree node)
depthFirstSearchFrom gr n = do
   MS.modify $ Set.delete n
   fmap (Tree.Node n . catMaybes) $
      Trav.for (Graph.successors gr n) $ \suc -> do
         unvisited <- MS.gets $ Set.member suc
         if unvisited
            then fmap Just $ depthFirstSearchFrom gr suc
            else pure Nothing



transposeMap :: (Ord k, Ord a) => Map k a -> Map a (NonEmptySet.T k)
transposeMap =
   Map.foldl (Map.unionWith NonEmptySet.union) Map.empty .
   Map.mapWithKey
      (\n root -> Map.singleton root $ NonEmptySet.singleton n)

{-
Direct imperative implementation of
<https://en.wikipedia.org/wiki/Kosaraju%27s_algorithm>
-}
stronglyConnectedComponents ::
   (Ord node) =>
   Graph Graph.DirEdge node edgeLabel nodeLabel -> [NonEmptySet.T node]
stronglyConnectedComponents gr =
   let visit u = do
         unvisited <- MS.gets snd
         when (Set.member u unvisited) $ do
            MS.modify (mapSnd $ Set.delete u)
            mapM_ visit $ Graph.successors gr u
            MS.modify (mapFst (u :))
       (queue, []) =
         mapSnd Set.toList $
         MS.execState
            (mapM_ visit $ Graph.nodes gr)
            ([], Graph.nodeSet gr)
       assign root u = do
         assignments <- MS.get
         when (Map.notMember u assignments) $ do
            MS.put $ Map.insert u root assignments
            mapM_ (assign root) $ Graph.predecessors gr u
   in Map.elems $ transposeMap $
      MS.execState (Fold.mapM_ (\n -> assign n n) queue) Map.empty