{-# language ViewPatterns #-}
module Agda.Utils.Graph.TopSort
( topSort
) where
import Data.List
import Data.Maybe
import Data.Function
import qualified Data.Set as Set
import qualified Data.Map as Map
import Control.Arrow
mergeBy :: (a -> a -> Bool) -> [a] -> [a] -> [a]
mergeBy _ [] xs = xs
mergeBy _ xs [] = xs
mergeBy f (x:xs) (y:ys)
| f x y = x: mergeBy f xs (y:ys)
| otherwise = y: mergeBy f (x:xs) ys
-- | topoligical sort with smallest-numbered available vertex first
-- | input: nodes, edges
-- | output is Nothing if the graph is not a DAG
topSort :: Ord n => [n] -> [(n, n)] -> Maybe [n]
topSort nodes edges = mergeBy (<) nodes' <$> g m is
where
g m []
| Map.null m = Just []
| otherwise = Nothing
g m ((n, Set.toList -> cs): ns)
= (n:) <$> g m' (mergeBy ((<) `on` fst) ns [(c, s) | c<-cs, (0, s) <- maybeToList $ Map.lookup c m'])
where
m' = foldr (Map.adjust $ first (+(-1))) (Map.delete n m) cs
is = map (second snd) . filter ((==0) . fst . snd) $ Map.toList m
nodes' = Set.toList $ Set.fromList nodes `Set.difference` Set.fromList (concatMap (\(a,b)->[a,b]) edges)
m = foldr f mempty $ nub edges
f (b, a)
= Map.alter (Just . maybe (1, mempty) (first (+1))) b
. Map.alter (Just . maybe (0, Set.singleton b) (second $ Set.insert b)) a