------------------------------------------------------------------------
-- Graphs represented as adjacency lists, along with some algorithms.
-- Author: Nils Anders Danielsson.
--
-- The code is partly based on "Structuring Depth-First Search
-- Algorithms in Haskell" by David J. King and John Launchbury
-- (http://doi.org/10.1145/199448.199530).
--
-- In the documentation of time complexity below, "v" stands for the
-- number of nodes, and "e" for the number of edges.
------------------------------------------------------------------------

{-# 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.List hiding (cycle)
import Data.Maybe
import Data.STRef.Lazy
import Prelude hiding (cycle)
import System.Directory
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

-- 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.
--
-- Linear in the size of the tree.

nodes :: Tree -> [Vertex]
nodes = preorderLR

-- 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

-- The nodes in the forest, in left-to-right level-order, with one
-- level per list. No inner list is empty.
--
-- Linear in the size of the forest.

levelOrderLR :: Forest -> [[Vertex]]
levelOrderLR ts = forest ts []
  where
  forest :: Forest -> [[Vertex]] -> [[Vertex]]
  forest []       = id
  forest (t : ts) = tree t . forest ts

  tree :: Tree -> [[Vertex]] -> [[Vertex]]
  tree (Root v children) vss =
    (v : vs') : forest [ t | Regular t <- children ] vss'
    where
    (vs', vss') = case vss of
      []       -> ([], [])
      vs : vss -> (vs, vss)

------------------------------------------------------------------------
-- 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 = [ (u, v) | (u, vs) <- assocs g, v <- vs ]

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

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

------------------------------------------------------------------------
-- Graph construction

-- 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 (reverse 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 ]

-- A single cycle with n nodes.
--
-- Precondition: n >= 1.
--
-- O(n).

cycle :: Integer -> Graph
cycle n = fromList (1, n) ((n, 1) : [ (i, i + 1) | i <- [1..n-1] ])

-- A "linear" graph with n nodes.
--
-- Precondition: n >= 0.
--
-- O(n).

line :: Integer -> Graph
line n = fromList (1, n) [ (i, i + 1) | i <- [1..n-1] ]

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

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

-- Merges two graphs, without adding any new edges. The nodes of the
-- second graph are (perhaps) relabelled, to avoid clashes with nodes
-- in the first graph.
--
-- O(v + e).

merge :: Graph -> Graph -> Graph
merge g1 g2 =
  fromList
    (from1, renumber to2)
    (edges g1 ++ [ (renumber u, renumber v) | (u, v) <- edges g2 ])
  where
  (from1, to1) = bounds g1
  (from2, to2) = bounds g2

  renumber u = u - from2 + to1 + 1

------------------------------------------------------------------------
-- Depth-first search and related algorithms

-- 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` preorderRL (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 nodes (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, 2)
  , (0, 3)
  , (2, 3)
  , (3, 0)
  , (4, 5)
  , (4, 4)
  , (5, 4)
  ]

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

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

g4 :: Graph
g4 = g2 `merge` fromList (7, 9) [(7, 8) , (8, 9)]

g5 :: Graph
g5 = 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)
  ]

------------------------------------------------------------------------
-- Displaying graphs and forests

-- NOTE: The commands below require that the "dot" program is present
-- in the PATH. This program can be downloaded from
-- http://www.graphviz.org/.
--
-- By default the commands use "dot" with the "-Txlib" option. This
-- may not work on all systems. On other systems one can instead use
-- an SVG image viewer, by setting imageViewer to "Just program",
-- where program is the command name used to invoke the image viewer.
-- Be aware that the image file is deleted when this command
-- terminates, so if the command spawns a different program and then
-- terminates, then the image might not be displayed at all.

imageViewer = Nothing

-- Runs "dot" with the given string as input, and tries to
-- display the resulting image, if any.

dot :: String -> IO ()
dot s = do
  forkIO $ case imageViewer of
    Nothing -> do
      withTempFile "dot" $ \(f, h) -> do
        hPutStr h s
        hClose h
        callProcess "dot" ["-Txlib", f]
    Just imageViewer ->
      withTempFile "svg" $ \(f, h) -> do
        hPutStr h =<< readProcess "dot" ["-Tsvg"] s
        hClose h
        callProcess imageViewer [f]
  return ()
  where
  withTempFile suffix =
    bracket
      (openTempFile "." ("." ++ suffix))
      (\(f, h) -> hClose h >> removeFile f)

  callProcess prog args =
    bracket
      (createProcess (proc prog args))
      (\(_, _, _, p) -> waitForProcess p)
      (\_ -> return ())

-- Opens a window containing a representation of the graph.

drawG :: Graph -> IO ()
drawG g =
  dot $
    "digraph \"A graph\" {\n" ++
       unlines (map show (vertices g)) ++
       unlines [ show u ++ " -> " ++ show v | (u, v) <- edges g ] ++
    "}\n"

-- Opens a window containing a representation of the forest.

drawF :: Forest -> IO ()
drawF f =
  dot $
    "digraph \"A forest\" {\n" ++
      "graph [ rankdir = TB, newrank = true ]\n" ++
      dummy ++
      clusters ++
      levels ++
    "}\n"
  where
  dummy = unlines $
    "dummy [ style = invis ]" :
    [ "dummy -> " ++ show u ++ " [ style = invis ]" | Root u _ <- f ]

  levels = unlines (map level (levelOrderLR f))
    where
    level vs =
      "{ rank = same; " ++ intercalate " " (map show vs) ++ " }"

  clusters = concat (zipWith tree [1..] f)
    where
    tree n t = unlines $
      ["subgraph cluster_" ++ show n ++ " {"] ++
         ["graph [ style = rounded ]"] ++
         [intercalate " " (map show (nodes t))] ++
         edgesT t [] ++
      ["}"]

    edgesT :: Tree -> [String] -> [String]
    edgesT (Root u es) ss = edgesE u es ss

    edgesE :: Vertex -> [Edge] -> [String] -> [String]
    edgesE u []                          ss = ss
    edgesE u (Regular t@(Root v _) : es) ss = edge Nothing u v :
                                              edgesT t (edgesE u es ss)
    edgesE u (Other k v            : es) ss = edge (Just k) u v :
                                              edgesE u es ss

    edge :: Maybe Kind -> Vertex -> Vertex -> String
    edge k u v = show u ++ " -> " ++ show v ++ attributes
      where
      attributes = case k of
        Nothing -> ""
        Just k  ->
          " [ style = " ++ style ++
           ", color = " ++ colour ++
           ", constraint = false" ++
          " ]"
          where
          style = case k of
            Back -> "dashed"
            _    -> "dotted"

          colour = case k of
            Back    -> "red"
            Forward -> "purple"
            Cross   -> "green3"