Haskell BFS/DFS enumerators

Graph.hs

module Main where

import Control.Monad.Identity (runIdentity)
import Data.Map (Map)
import Data.Set (Set)
import Data.Maybe (isNothing)

import qualified Data.Map as Map
import qualified Data.Set as Set

import Data.Enumerator (
  Enumerator, Iteratee, Step(..), Stream(..),
  (>>==), ($$), ($=), (=$), run_, returnI, continue, yield
  )
import qualified Data.Enumerator.List as EL

----------------------------------------------------------------------

class GraphVisitContainer c where
  emptyContainer     :: c a
  nextNode           :: c a -> Maybe (Maybe a, a, c a)
  appendNodesToVisit :: (Ord a)
                     => Maybe a     -- parent
                     -> [a]         -- children
                     -> c a         -- container before adding
                     -> c a         -- container after adding

----------------------------------------------------------------------

type Graph a = Map a (Node a)
data Node a
  = Node {
    nodeValue    :: a
  , nodeIndex    :: Int
  , nodeSiblings :: [a]
  }
  deriving (Show, Eq)

data NodeVisit a
  = NodeVisit {
    visitingNode    :: Node a
  , fromNode        :: Maybe (Node a)
  , visitIndex      :: Int
  }
  deriving (Show, Eq)

--------------------

nodeVisited :: Node a -> Bool
nodeVisited = (>= 0) . nodeIndex

visitNode :: Ord a => Int -> a -> Graph a -> Graph a
visitNode index nodeVal graph = Map.adjust markVisited nodeVal graph
  where
    markVisited node = node { nodeIndex = index }

buildGraph :: (Ord a) => [(a, [a])] -> Graph a
buildGraph = Map.fromList . map toNode
  where
    toNode (val, siblings) = (val, Node val (-1) siblings)

--------------------

enumGraph :: (Monad m, GraphVisitContainer c, Ord a)
          => Int
          -> c a
          -> Graph a
          -> Enumerator (NodeVisit a) m b
enumGraph travIndex pendingVisit graph step@(Continue kn) =
    case nextNode pendingVisit of
      Nothing -> returnI step
      Just (parent, current, pendingVisit') -> do
        let parentNode = parent >>= (`Map.lookup` graph)
        case Map.lookup current graph of
          Nothing -> error "ERROR: trying to lookup non existent node in graph"
          Just currentNode
            | nodeVisited currentNode ->
                kn (Chunks [NodeVisit currentNode
                                           parentNode
                                           travIndex]) >>==
                enumGraph (travIndex + 1) pendingVisit' graph
            | otherwise ->
              let
                pendingVisit'' = appendNodesToVisit (Just current)
                                                    (nodeSiblings currentNode)
                                                    pendingVisit'
                graph' = visitNode travIndex current graph
              in
                kn (Chunks [NodeVisit currentNode
                                      parentNode
                                      travIndex]) >>==
                enumGraph (travIndex + 1) pendingVisit'' graph'
enumGraph _ _ _ step = returnI step

----------------------------------------------------------------------
-- BFS Logic

newtype BFSContainer a
  = BFSContainer (Int, Set (Int, Maybe a, a))
  deriving (Show, Eq)


instance GraphVisitContainer BFSContainer where
  emptyContainer = BFSContainer (0, Set.empty)

  nextNode (BFSContainer (bfsCount, nodeSet0)) =
      case Set.minView nodeSet0 of
        Nothing -> Nothing
        Just ((_, parent, current), nodeSet) ->
          Just (parent, current, BFSContainer (bfsCount, nodeSet))

  appendNodesToVisit parentNode
                     childrenNodes
                     (BFSContainer (bfsCount0, nodeSet0)) =
    let bfsCount = bfsCount0 + Set.size nodeSet
        nodeSet  = Set.union nodeSet0 .
                   Set.fromList $
                   zip3 [bfsCount0 + 1..]
                        (repeat parentNode)
                        childrenNodes
    in BFSContainer (bfsCount, nodeSet)

--------------------

enumBFS :: (Monad m, Ord a)
        => [a]
        -> Graph a
        -> Enumerator (NodeVisit a) m b
enumBFS nodeList initialGraph =
    enumGraph 0 container initialGraph
  where
    container = appendNodesToVisit Nothing
                                   nodeList
                                   (emptyContainer :: BFSContainer a)


----------------------------------------------------------------------
-- DFS Logic

newtype DFSContainer a
  = DFSContainer (Int, Set (Int, Maybe a, a))
  deriving (Show, Eq)


instance GraphVisitContainer DFSContainer where
  emptyContainer = DFSContainer (0, Set.empty)

  nextNode (DFSContainer (dfsCount, nodeSet0)) =
      case Set.maxView nodeSet0 of
        Nothing -> Nothing
        Just ((_, parent, current), nodeSet) ->
          Just (parent, current, DFSContainer (dfsCount, nodeSet))

  appendNodesToVisit parentNode
                     childrenNodes
                     (DFSContainer (dfsCount0, nodeSet0)) =
    let dfsCount = dfsCount0 + 1
        nodeSet  = Set.union nodeSet0 .
                   Set.fromList $
                   zip3 (repeat dfsCount)
                        (repeat parentNode)
                        childrenNodes
    in DFSContainer (dfsCount, nodeSet)

--------------------

enumDFS :: (Monad m, Ord a)
        => [a]
        -> Graph a
        -> Enumerator (NodeVisit a) m b
enumDFS nodeList initialGraph =
    enumGraph 0 container initialGraph
  where
    container = appendNodesToVisit Nothing
                                   nodeList
                                   (emptyContainer :: DFSContainer a)

----------------------------------------------------------------------
-- Bipartite Graphs

bipartiteIteratee :: (Ord a, Monad m) => Iteratee (NodeVisit a) m Bool
bipartiteIteratee = continue go
  where
    go EOF = yield True EOF
    go (Chunks [visit])
      | (nodeVisited $ visitingNode visit) =
        case fromNode visit of
          Just parentNode
            | nodeIndex parentNode `mod` 2 ==
              nodeIndex (visitingNode visit) `mod` 2 -> yield False (Chunks [])
            | otherwise -> continue go
          Nothing -> continue go
      | otherwise = continue go

graphIsBipartite :: Ord a => Graph a -> Bool
graphIsBipartite g = 
    runIdentity (run_ $ enumDFS (Map.keys g) g $$ bipartiteIteratee)

----------------------------------------------------------------------
-- Connected Graph

isConnectedIteratee :: (Ord a, Monad m) => Iteratee (NodeVisit a) m Bool
isConnectedIteratee = continue $ go 0
  where
    go discCount EOF = yield True EOF
    go discCount (Chunks [visit])
      | isNothing (fromNode visit) && 
        (discCount + 1) > 1 = yield False (Chunks [])
      | otherwise = continue $ go (discCount + 1)

isConnectedGraph :: Ord a => Graph a -> Bool
isConnectedGraph g = 
    runIdentity (run_ $ enumDFS (Map.keys g) g $$ isConnectedIteratee)

----------------------------------------------------------------------

main :: IO ()
main = do
    putStrLn "DFS exec:"
    print $ graphIsBipartite graph
    print $ isConnectedGraph graph
    run_ (enumDFS "a" graph $$ EL.consume) >>= mapM_ print
    -- putStrLn "BFS exec:"
    -- run_ (enumBFS "a" graph $$ EL.consume) >>= mapM_ print
  where
    graph = buildGraph [('a', "bcd"), ('b', "acf"), ('c', "adf"), ('d', "acf"), ('f', "bcd")]