------------------------------------------------------------------------
-- Graphs represented as adjacency lists, along with some algorithms
--
-- Nils Anders Danielsson
------------------------------------------------------------------------

-- The code is partly based on "Structuring Depth-First Search
-- Algorithms in Haskell" by David J. King and John Launchbury.

{-# LANGUAGE ScopedTypeVariables #-}

module AdjacencyList where

import Control.Applicative
import Control.Concurrent
import Control.Exception
import Control.Monad.ST.Lazy
import Data.Array.IArray
import Data.Array.ST
import Data.Char
import Data.Maybe
import Data.STRef.Lazy
import qualified Data.Tree as T
import qualified Data.Tree.Pretty as T
import System.IO
import System.Process

------------------------------------------------------------------------
-- Depth-first spanning trees and forests

-- Depth-first spanning forests.

type Forest = [Tree]

-- Depth-first spanning trees.

data Tree = Root Vertex [Edge]
  deriving Show

-- Depth-first spanning tree edges.

data Edge = Regular Tree
          | Other Kind Vertex
  deriving Show

-- The different kinds of non-tree edges.

data Kind = Forward | Back | Cross
  deriving Show

-- Pretty-prints a depth-first spanning forest.

pretty :: Forest -> IO ()
pretty = putStr . T.drawVerticalForest . map toTreeT
  where
  toTreeT :: Tree -> T.Tree String
  toTreeT (Root v es) = T.Node (show v) (map toTreeE es)

  toTreeE :: Edge -> T.Tree String
  toTreeE (Regular t) = toTreeT t
  toTreeE (Other k v) =
    T.Node (map toLower (show k) ++ ": " ++ show v) []

-- Does the forest contain a back edge?
--
-- Linear in the size of the forest.

containsBackEdge :: Forest -> Bool
containsBackEdge ts = any containsBackEdgeT ts
  where
  containsBackEdgeT :: Tree -> Bool
  containsBackEdgeT (Root _ children) = any containsBackEdgeE children

  containsBackEdgeE :: Edge -> Bool
  containsBackEdgeE (Regular t)    = containsBackEdgeT t
  containsBackEdgeE (Other Back _) = True
  containsBackEdgeE _              = False

-- The nodes in the tree, in left-to-right preorder.
--
-- Linear in the size of the tree.

preorderLR :: Tree -> [Vertex]
preorderLR t = tree t []
  where
  tree :: Tree -> [Vertex] -> [Vertex]
  tree (Root v children) = (v :) . edges children

  edges :: [Edge] -> [Vertex] -> [Vertex]
  edges []               = id
  edges (Regular t : es) = tree t . edges es
  edges (_         : es) = edges es

-- The nodes in the forest, in right-to-left preorder.
--
-- Note that the implementation can be seen as performing a
-- /left-to-right postorder/ traversal of the forest, and pushing
-- nodes onto an initially empty stack.
--
-- Linear in the size of the forest.

preorderRL :: Forest -> [Vertex]
preorderRL ts = trees ts []
  where
  trees :: [Tree] -> [Vertex] -> [Vertex]
  trees []                           = id
  trees (Root v children : siblings) =
    trees siblings . (v :) . edges children

  edges :: [Edge] -> [Vertex] -> [Vertex]
  edges []               = id
  edges (Regular t : es) = edges es . trees [t]
  edges (_         : es) = edges es

------------------------------------------------------------------------
-- Graphs

-- Vertices. For simplicity just integers.

type Vertex = Integer

-- Directed, unweighted graphs represented using adjacency lists.

type Graph = Array Vertex [Vertex]

-- Returns the graph's vertices (in ascending order).
--
-- O(v).

vertices :: Graph -> [Vertex]
vertices g = indices g

-- Returns the graph's edges. The tuples have the form (from, to).
--
-- O(v + e).

edges :: Graph -> [(Vertex, Vertex)]
edges g =
  concat $
  map (\(u, vs) -> map (\v -> (u, v)) vs) $
  assocs g

-- Tries to display the graph represented by the given string, which
-- should be written in the DOT language, on the screen.
--
-- The command attempts to use the dot program, with the -Txlib
-- option. This is not guaranteed to work on all systems.

dot :: String -> IO ()
dot s =
  bracket
    (createProcess ((proc "dot" ["-Txlib"]) { std_in = CreatePipe }))
    (\(Just h, _, _, p) ->
        hClose h
          `finally`
        -- Avoid zombies.
        forkIO (waitForProcess p >> return ()))
    (\(Just h, _, _, _) -> hPutStr h s)

-- Opens a window containing a representation of the graph.
--
-- The command attempts to use the dot program, with the -Txlib
-- option. This is not guaranteed to work on all systems.

draw :: Graph -> IO ()
draw g =
  dot $
    "digraph G {\n" ++
    unlines (map show (vertices g)) ++
    unlines (map (\(u, v) -> show u ++ " -> " ++ show v) (edges g)) ++
    "}\n"

-- Returns the nodes adjacent to the given node.
--
-- O(1).

adjacent :: Vertex -> Graph -> [Vertex]
adjacent v g = g ! v

-- Constructs a graph from the given node range and the given list of
-- edges. The tuples have the form (from, to).
--
-- Precondition: Every node in the list of edges must be in range.
-- Precondition: No edge may occur twice in the list.
--
-- O(v + e).

fromList :: (Vertex, Vertex) -> [(Vertex, Vertex)] -> Graph
fromList bounds es = accumArray (flip (:)) [] bounds es

-- Complete graph with n nodes.
--
-- Precondition: n >= 0.
--
-- O(n²).

complete :: Integer -> Graph
complete n =
  fromList (1, n) [(i, j) | i <- [1..n], j <- [1..n], i /= j ]

-- Constructs the opposite graph.
--
-- O(v + e).

opposite :: Graph -> Graph
opposite g = fromList (bounds g) [ (v, u) | (u, v) <- edges g ]

-- Depth-first search.
--
-- The list of nodes determines the order of searches: every search
-- starts from the earliest unvisited node. Only nodes reachable from
-- nodes in the list will be visited.
--
-- Precondition: Every node in the list must exist in the graph.
--
-- O(v + e).

dfs :: Graph -> [Vertex] -> Forest
dfs g vs = runST $ do
  -- Nodes are numbered sequentially in the order in which they are
  -- visited. This counter keeps track of the number of the next node.
  (counter :: STRef s Integer) <- newSTRef 0

  -- Just <sequence number> for visited nodes, Nothing for others.
  (visited :: STArray s Vertex (Maybe Integer))
    <- newArray (bounds g) Nothing

  -- True for active nodes: nodes that we have started visiting, but
  -- not yet finished processing.
  (active :: STArray s Vertex Bool) <- newArray (bounds g) False

  let
    -- Is the given node visited? If so, its sequence number is
    -- returned.
    isVisited :: Vertex -> ST s (Maybe Integer)
    isVisited v = readArray visited v

    -- Marks a node as visited. The sequence number is returned.
    markVisited :: Vertex -> ST s Integer
    markVisited v = do
      c <- readSTRef counter
      writeSTRef counter (1 + c)
      writeArray visited v (Just c)
      return c

    -- Is the given node active?
    isActive :: Vertex -> ST s Bool
    isActive v = readArray active v

    -- Run the given computation with the given node activated.
    activate :: Vertex -> ST s a -> ST s a
    activate v c = do
      writeArray active v True
      x <- c
      writeArray active v False
      return x

    -- Computes a tree for the given vertex.
    -- Precondition: The vertex must not have been visited.
    dfsVertex :: Vertex -> ST s Tree
    dfsVertex v = do
      n <- markVisited v
      activate v $
        Root v <$> mapM (dfsEdge n) (adjacent v g)

    -- dfsEdge n v computes a tree for the edge going from the node
    -- with sequence number n to v.
    dfsEdge :: Integer -> Vertex -> ST s Edge
    dfsEdge parentNumber v = do
      visited <- isVisited v
      case visited of
        Nothing      -> Regular <$> dfsVertex v
        Just vNumber -> do
          active <- isActive v
          return $
            if active then
              Other Back v
            else if parentNumber > vNumber then
              Other Cross v
            else if parentNumber < vNumber then
              Other Forward v
            else
              error "dfs: Impossible case."

  -- Computes the forest.
  catMaybes <$>
    mapM (\v -> do
            visited <- isVisited v
            case visited of
              Just _  -> return Nothing
              Nothing -> Just <$> dfsVertex v)
         vs

-- Computes a depth-first spanning forest for the graph.
--
-- O(v + e).

dff :: Graph -> Forest
dff g = dfs g (vertices g)

-- "reachable g from to" is True iff "to" can be reached from "from"
-- in "g".
--
-- Precondition: "from" must be present in "g".
--
-- O(v + e).

reachable :: Graph -> Vertex -> Vertex -> Bool
reachable g from to =
  to `elem` concat (map preorderLR (dfs g [from]))

-- Is the graph cyclic?
--
-- O(e + v).

cyclic :: Graph -> Bool
cyclic g = containsBackEdge (dff g)

-- Returns Nothing if the graph is cyclic, and otherwise a
-- topologically sorted list of the nodes in the graph.
--
-- O(e + v).

topologicalSort :: Graph -> Maybe [Vertex]
topologicalSort g
  | containsBackEdge forest = Nothing
  | otherwise               = Just (preorderRL forest)
  where forest = dff g

-- Strongly connected components.

type SCC = [Vertex]

-- Computes strongly connected components. The SCC list is sorted in
-- reverse topological order.
--
-- O(e + v).

sccs :: Graph -> [SCC]
sccs g = map preorderLR (dfs g (preorderRL (dff (opposite g))))

-- Is the graph strongly connected?
--
-- O(e + v).

stronglyConnected :: Graph -> Bool
stronglyConnected g = case sccs g of
  _ : _ : _ -> False
  _         -> True

------------------------------------------------------------------------
-- Examples

g1 :: Graph
g1 = fromList (0, 5)
  [ (1, 0)
  , (0, 3)
  , (0, 2)
  , (2, 3)
  , (3, 0)
  , (4, 4)
  , (4, 5)
  , (5, 4)
  ]

g2 :: Graph
g2 = fromList (1, 6)
  [ (1, 3)
  , (1, 5)
  , (2, 1)
  , (2, 5)
  , (3, 1)
  , (3, 3)
  , (3, 4)
  , (3, 6)
  , (4, 1)
  , (5, 6)
  , (6, 5)
  , (6, 6)
  ]

g3 :: Graph
g3 = opposite g2

g4 :: Graph
g4 = fromList (1, 7)
  [ (1, 3)
  , (2, 3)
  , (2, 4)
  , (3, 5)
  , (3, 6)
  , (4, 6)
  , (5, 7)
  , (6, 7)
  ]