content: -- A small library of functions on directed graphs
  -- using a simple list-of-successors representation.
  -- Colin Runciman, May 2015
  
  module Digraph (Digraph(..), okDigraph, strictOrder,
                  sources, targets, nodes, preds, succs,
                  isNode, isEdge, isPath,
                  emptyDigraph, addNode, addEdge, assoc1toNdigraph,
                  transitiveClosure, topoSort,
                  insert, union, diff, cycles, subgraph, maxDagFrom) where
  
  import GHC.Exts (groupWith)
  import Data.List (partition,(\\),sort)
  import Data.Maybe (isJust,fromJust)
  import Control.Monad (guard)
  
  data Digraph a = D {nodeSuccs :: [(a,[a])]} deriving (Eq, Show)
  -- Data invariant: in a digraph pair-list [...(source,targets)...]:
  -- (1) pairs are listed in strictly increasing source order
  -- (2) each list of targets is in strictly increasing order
  -- (3) every element in a list of targets must itself be
  --     listed as a "source", though perhaps with [] targets
  
  okDigraph :: (Ord a, Eq a) => Digraph a -> Bool
  okDigraph (D d)  =
    strictOrder ss && all goodTargetList tss
    where
    ss                     =  map fst d
    tss                    =  map snd d
    goodTargetList ts      =  strictOrder ts &&
                              all (`elemOrd` ss) ts
  
  strictOrder :: (Ord a, Eq a) => [a] -> Bool
  strictOrder (x:y:etc)  =  x < y && strictOrder (y:etc)
  strictOrder _          =  True
  
  nodes :: Digraph a -> [a]
  nodes (D d)  =  [s | (s,_) <- d]
  
  sources :: Digraph a -> [a]
  sources (D d)  =  [s | (s,ts) <- d, not (null ts)]
  
  targets :: (Ord a, Eq a) => Digraph a -> [a]
  targets (D d) =  foldr union [] (map snd d)
  
  preds :: (Ord a, Eq a) => a -> Digraph a -> [a]
  preds t (D d)  =  [s | (s,ts) <- d, t `elemOrd` ts]
  
  succs :: (Ord a, Eq a) => a -> Digraph a -> [a]
  succs s (D d)  =  case lookup s d of
                    Just ns -> ns
                    Nothing -> []
  
  isNode :: (Ord a, Eq a) => a -> Digraph a -> Bool
  isNode n (D d)  =  isJust (lookup n d)
  
  isEdge :: (Ord a, Eq a) => a -> a -> Digraph a -> Bool
  isEdge s t (D d)  =  case lookup s d of
                       Just ns -> t `elemOrd` ns
                       Nothing -> False
  
  isPath :: (Ord a, Eq a) => a -> a -> Digraph a -> Bool
  isPath s t d  =  t `elemOrd` closeInto d [] [s]
  
  emptyDigraph :: Digraph a
  emptyDigraph  =  D []
  
  addNode :: (Ord a, Eq a) => a -> Digraph a -> Digraph a
  addNode s (D d)  =
    let (these,those)  =  span ((< s) . fst) d in
    D $ these ++
    case those of
    []            ->  [(s,[])]
    (sd,tsd):etc  ->  if s == sd then error "addNode: already present"
                      else (s,[]) : those
  
  addEdge :: (Ord a, Eq a) => a -> a -> Digraph a -> Digraph a
  addEdge s t (D d)  =
    let (these,those)  =  span ((< s) . fst) d in
    D $ these ++
    case those of
    []            ->  [(s,[t])]
    (sd,tsd):etc  ->  if s == sd then
                         if t `elemOrd` tsd then error "addEdge: already present"
                         else (s,insert t tsd) : etc
                      else (s,[t]) : those
  
  -- The function assoc1toNdigraph derives a digraph from an association list
  -- pairing single sources with lists of targets.  Sorting is applied to
  -- outer and inner lists, so there is no ordering requirement on the argument.
  -- If there is more than one pair with the same source, target lists are merged.
  -- If the same value appears more than once in a target list, duplicates are removed.
  -- If any target does not occur as a source, it is added, with an empty target list.
  assoc1toNdigraph :: (Ord a, Eq a) => [(a,[a])] -> Digraph a
  assoc1toNdigraph stss  =  D $ addMissingSources $ mergeAndSortTargets $ sort stss
    where
    mergeAndSortTargets []                       =  []
    mergeAndSortTargets [(s,ts)]                 =  [(s, nubOrd $ sort ts)]
    mergeAndSortTargets ((s0,ts0):(s1,ts1):etc)  =
      if s0 == s1 then mergeAndSortTargets ((s0,ts0++ts1):etc)
      else (s0,nubOrd $ sort ts0) : mergeAndSortTargets ((s1,ts1):etc)
    addMissingSources stss  =
      union [(s,[]) | s <- missingSources] stss
      where
      missingSources  =  allTargets `diff` map fst stss
      allTargets      =  foldr union [] (map snd stss)
  
  transitiveClosure :: (Ord a, Eq a) => Digraph a -> Digraph a
  transitiveClosure d =  D $ map (close d) (nodeSuccs d)
  
  close :: (Ord a, Eq a) => Digraph a -> (a,[a]) -> (a,[a])
  close d (s,ts)  =  (s, closeInto d [] ts)
  
  closeInto :: (Ord a, Eq a) => Digraph a -> [a] -> [a] -> [a]
  closeInto d clo []      =  clo
  closeInto d clo (t:ts)  =
    case lookup t (nodeSuccs d) of
    Nothing     ->  closeInto d clo ts
    Just tsuccs ->  closeInto d clo' (union (diff tsuccs clo') ts)
    where
    clo'  =  insert t clo
  
  -- auxiliary functions for ordered list processing
  
  nubOrd :: (Ord a, Eq a) => [a] -> [a]
  nubOrd []         =  []
  nubOrd [x]        =  [x]
  nubOrd (x:y:etc)  =  if x==y then nubOrd (y:etc) else x : nubOrd (y:etc)
  
  elemOrd :: Ord a => a -> [a] -> Bool
  elemOrd x ys  =  null (diff [x] ys)
  
  insert :: Ord a => a -> [a] -> [a]
  insert x ys  =  union [x] ys
  
  union :: Ord a => [a] -> [a] -> [a]
  union [] ys         =  ys
  union (x:xs) []     =  x:xs
  union (x:xs) (y:ys) =  case compare x y of
                         LT -> x : union xs (y:ys)
                         EQ -> union xs (y:ys)
                         GT -> y : union (x:xs) ys
  
  diff :: Ord a => [a] -> [a] -> [a]
  diff [] ys          =  []
  diff (x:xs) []      =  x:xs
  diff (x:xs) (y:ys)  =  case compare x y of
                         LT -> x : diff xs (y:ys)
                         EQ -> diff xs ys
                         GT -> diff (x:xs) ys
  
  -- The result of cycles lists disjoint maximal subsets of nodes in
  -- each of which there is a cycle passing through all nodes.
  cycles:: (Ord a, Eq a) => Digraph a -> [[a]]
  cycles d  =
    let d'          =  transitiveClosure d
        cycleNodes  =  filter (hasLoop d') (sources d')
    in
    map (sources . D) $ groupWith snd $ nodeSuccs $ subgraph cycleNodes d'
  
  hasLoop :: Eq a => Digraph a -> a -> Bool
  hasLoop d s  =
    case lookup s (nodeSuccs d) of
    Nothing -> False
    Just ts -> elem s ts
  
  subgraph :: Eq a => [a] -> Digraph a -> Digraph a
  subgraph ns d  =  D $ [(s,filter (`elem` ns) ts) | (s,ts) <- nodeSuccs d, elem s ns]
  
  -- The result of topoSort d, where d is an acyclic digraph, lists all nodes
  -- in an order where each node precedes all its digraph successors;
  -- the result is Nothing if d has a cycle.
  topoSort :: (Ord a, Eq a) => Digraph a -> Maybe [a]
  topoSort (D [])  =  Just []
  topoSort d       =  do
                         guard (not $ null maxima)
                         ns <- topoSort (D nodesuccs')
                         return $ ns ++ maxima
    where
    (these,those)  =  partition (null.snd) (nodeSuccs d)
    maxima         =  nodes (D these)
    nodesuccs'     =  [(s,ts \\ maxima) | (s,ts) <- those]
  
  -- The result of maxDAGfrom s d is a subgraph of d which is a maximal
  -- DAG rooted at node s.
  maxDagFrom :: (Ord a, Eq a) => a -> Digraph a -> Digraph a
  maxDagFrom s d  =  md [] [s] (removeAllEdges d) (removeLoops d)
  
  -- In a call md done todo dag d, done is an ordered list of nodes already
  -- visited, todo is a disjoint ordered list of nodes to be visited, dag is
  -- the dag so far, and d is the full loop-free digraph.
  md :: (Ord a, Eq a) => [a] -> [a] -> Digraph a -> Digraph a -> Digraph a
  md _    []     dag d  =  dag
  md done (s:ss) dag d  =
    case lookup s (nodeSuccs d) of
    Nothing  ->  md done' ss dag  d
    Just ts  ->  md done' (union (diff ts done) ss) dag' d
                 where
                 dag'  =  foldr (uncurry addEdgeIfAcyclic) dag [(s,t) | t <- ts]
    where
    done'  =  insert s done
  
  addEdgeIfAcyclic :: (Ord a, Eq a) => a -> a -> Digraph a -> Digraph a
  addEdgeIfAcyclic s t d  =  if isPath t s d then d else addEdge s t d
  
  removeLoops :: (Ord a, Eq a) => Digraph a -> Digraph a
  removeLoops d  =  D $ [(s, diff ts [s]) | (s, ts) <- nodeSuccs d]
  
  removeAllEdges :: (Ord a, Eq a) => Digraph a -> Digraph a
  removeAllEdges d  =  D $ [(s, []) | (s, _) <- nodeSuccs d]
  
  