-- | Utilities for sorting

module Hydra.Tools.Sorting (
  topologicalSort,
  topologicalSortComponents,
  initGraph
) where

import qualified Data.List as L
import qualified Data.Bifunctor as BF

import qualified Data.Graph as G


-- | Sort a directed acyclic graph (DAG) based on an adjacency list
--   Yields a list of nontrivial strongly connected components if the graph has cycles
topologicalSort :: Ord a => [(a, [a])] -> Either [[a]] [a]
topologicalSort :: forall a. Ord a => [(a, [a])] -> Either [[a]] [a]
topologicalSort [(a, [a])]
pairs = if [[a]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
L.null [[a]]
withCycles
    then [a] -> Either [[a]] [a]
forall a b. b -> Either a b
Right ([a] -> Either [[a]] [a]) -> [a] -> Either [[a]] [a]
forall a b. (a -> b) -> a -> b
$ [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat [[a]]
sccs
    else [[a]] -> Either [[a]] [a]
forall a b. a -> Either a b
Left [[a]]
withCycles
  where
    sccs :: [[a]]
sccs = [(a, [a])] -> [[a]]
forall a. Ord a => [(a, [a])] -> [[a]]
topologicalSortComponents [(a, [a])]
pairs
    withCycles :: [[a]]
withCycles = ([a] -> Bool) -> [[a]] -> [[a]]
forall a. (a -> Bool) -> [a] -> [a]
L.filter [a] -> Bool
forall a. [a] -> Bool
isCycle [[a]]
sccs
    isCycle :: [a] -> Bool
isCycle = Bool -> Bool
not (Bool -> Bool) -> ([a] -> Bool) -> [a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
L.null ([a] -> Bool) -> ([a] -> [a]) -> [a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. HasCallStack => [a] -> [a]
L.tail

-- | Find the strongly connected components (including cycles and isolated vertices) of a graph, in (reverse) topological order
topologicalSortComponents :: Ord a => [(a, [a])] -> [[a]]
topologicalSortComponents :: forall a. Ord a => [(a, [a])] -> [[a]]
topologicalSortComponents [(a, [a])]
pairs = [a] -> [a]
forall a. Ord a => [a] -> [a]
L.sort ([a] -> [a]) -> (Tree Vertex -> [a]) -> Tree Vertex -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Vertex -> a) -> [Vertex] -> [a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Vertex -> ((), a, [a])) -> Vertex -> a
forall a. (Vertex -> ((), a, [a])) -> Vertex -> a
getKey Vertex -> ((), a, [a])
nodeFromVertex) ([Vertex] -> [a])
-> (Tree Vertex -> [Vertex]) -> Tree Vertex -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree Vertex -> [Vertex]
forall a. Tree a -> [a]
treeToList) (Tree Vertex -> [a]) -> [Tree Vertex] -> [[a]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Graph -> [Tree Vertex]
G.scc Graph
g
  where
    (Graph
g, Vertex -> ((), a, [a])
nodeFromVertex) = [(a, [a])] -> (Graph, Vertex -> ((), a, [a]))
forall a. Ord a => [(a, [a])] -> (Graph, Vertex -> ((), a, [a]))
initGraph [(a, [a])]
pairs

treeToList :: G.Tree a -> [a]
treeToList :: forall a. Tree a -> [a]
treeToList (G.Node a
root [Tree a]
subforest) = a
roota -> [a] -> [a]
forall a. a -> [a] -> [a]
:([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat (Tree a -> [a]
forall a. Tree a -> [a]
treeToList (Tree a -> [a]) -> [Tree a] -> [[a]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Tree a]
subforest))

getKey :: (G.Vertex -> ((), a, [a])) -> G.Vertex -> a
getKey :: forall a. (Vertex -> ((), a, [a])) -> Vertex -> a
getKey Vertex -> ((), a, [a])
nodeFromVertex Vertex
v = a
n
  where
    (()
_, a
n, [a]
_) = Vertex -> ((), a, [a])
nodeFromVertex Vertex
v

initGraph :: Ord a => [(a, [a])] -> (G.Graph, G.Vertex -> ((), a, [a]))
initGraph :: forall a. Ord a => [(a, [a])] -> (Graph, Vertex -> ((), a, [a]))
initGraph [(a, [a])]
pairs = (Graph
g, Vertex -> ((), a, [a])
nodeFromVertex)
  where
    (Graph
g, Vertex -> ((), a, [a])
nodeFromVertex, a -> Maybe Vertex
_) = [((), a, [a])]
-> (Graph, Vertex -> ((), a, [a]), a -> Maybe Vertex)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)
G.graphFromEdges ((a, [a]) -> ((), a, [a])
forall {b} {c}. (b, c) -> ((), b, c)
toEdge ((a, [a]) -> ((), a, [a])) -> [(a, [a])] -> [((), a, [a])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, [a])]
pairs)
    toEdge :: (b, c) -> ((), b, c)
toEdge (b
key, c
keys) = ((), b
key, c
keys)