module Agda.Utils.Graph.TopSort
    ( topSort
    ) where

import qualified Data.Set as Set
import qualified Data.Map as Map
import qualified Agda.Utils.Graph.AdjacencyMap.Unidirectional as G

-- NB:: Defined but not used
mergeBy :: (a -> a -> Bool) -> [a] -> [a] -> [a]
mergeBy :: forall a. (a -> a -> Bool) -> [a] -> [a] -> [a]
mergeBy a -> a -> Bool
_ [] [a]
xs = [a]
xs
mergeBy a -> a -> Bool
_ [a]
xs [] = [a]
xs
mergeBy a -> a -> Bool
f (a
x:[a]
xs) (a
y:[a]
ys)
    | a -> a -> Bool
f a
x a
y = a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a -> Bool) -> [a] -> [a] -> [a]
forall a. (a -> a -> Bool) -> [a] -> [a] -> [a]
mergeBy a -> a -> Bool
f [a]
xs (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys)
    | Bool
otherwise = a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a -> Bool) -> [a] -> [a] -> [a]
forall a. (a -> a -> Bool) -> [a] -> [a] -> [a]
mergeBy a -> a -> Bool
f (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs) [a]
ys

-- | topoligical sort with smallest-numbered available vertex first
-- | input: nodes, edges
-- | output is Nothing if the graph is not a DAG
--   Note: should be stable to preserve order of generalizable variables. Algorithm due to Richard
--   Eisenberg, and works by walking over the list left-to-right and moving each node the minimum
--   distance left to guarantee topological ordering.
topSort :: Ord n => [n] -> [(n, n)] -> Maybe [n]
topSort :: forall n. Ord n => [n] -> [(n, n)] -> Maybe [n]
topSort [n]
nodes [(n, n)]
edges = [(n, Set n)] -> [n] -> Maybe [n]
go [] [n]
nodes
  where
    -- #4253: The input edges do not necessarily include transitive dependencies, so take transitive
    --        closure before sorting.
    w :: Maybe ()
w      = () -> Maybe ()
forall a. a -> Maybe a
Just () -- () is not a good edge label since it counts as a "zero" edge and will be ignored
    g :: Graph n (Maybe ())
g      = Graph n (Maybe ()) -> Graph n (Maybe ())
forall n e. (Ord n, Eq e, StarSemiRing e) => Graph n e -> Graph n e
G.transitiveClosure (Graph n (Maybe ()) -> Graph n (Maybe ()))
-> Graph n (Maybe ()) -> Graph n (Maybe ())
forall a b. (a -> b) -> a -> b
$ [n] -> Graph n (Maybe ())
forall n e. Ord n => [n] -> Graph n e
G.fromNodes [n]
nodes Graph n (Maybe ()) -> Graph n (Maybe ()) -> Graph n (Maybe ())
forall n e. Ord n => Graph n e -> Graph n e -> Graph n e
`G.union` [Edge n (Maybe ())] -> Graph n (Maybe ())
forall n e. Ord n => [Edge n e] -> Graph n e
G.fromEdges [n -> n -> Maybe () -> Edge n (Maybe ())
forall n e. n -> n -> e -> Edge n e
G.Edge n
a n
b Maybe ()
w | (n
a, n
b) <- [(n, n)]
edges]
    deps :: n -> Set n
deps n
a = Map n (Maybe ()) -> Set n
forall k a. Map k a -> Set k
Map.keysSet (Map n (Maybe ()) -> Set n) -> Map n (Maybe ()) -> Set n
forall a b. (a -> b) -> a -> b
$ Graph n (Maybe ()) -> Map n (Map n (Maybe ()))
forall n e. Graph n e -> Map n (Map n e)
G.graph Graph n (Maybe ())
g Map n (Map n (Maybe ())) -> n -> Map n (Maybe ())
forall k a. Ord k => Map k a -> k -> a
Map.! n
a

    -- acc: Already sorted nodes in reverse order paired with accumulated set of nodes that must
    -- come before it
    go :: [(n, Set n)] -> [n] -> Maybe [n]
go [(n, Set n)]
acc [] = [n] -> Maybe [n]
forall a. a -> Maybe a
Just ([n] -> Maybe [n]) -> [n] -> Maybe [n]
forall a b. (a -> b) -> a -> b
$ [n] -> [n]
forall a. [a] -> [a]
reverse ([n] -> [n]) -> [n] -> [n]
forall a b. (a -> b) -> a -> b
$ ((n, Set n) -> n) -> [(n, Set n)] -> [n]
forall a b. (a -> b) -> [a] -> [b]
map (n, Set n) -> n
forall a b. (a, b) -> a
fst [(n, Set n)]
acc
    go [(n, Set n)]
acc (n
n : [n]
ns) = ([(n, Set n)] -> [n] -> Maybe [n]
`go` [n]
ns) ([(n, Set n)] -> Maybe [n]) -> Maybe [(n, Set n)] -> Maybe [n]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< n -> [(n, Set n)] -> Maybe [(n, Set n)]
insert n
n [(n, Set n)]
acc

    insert :: n -> [(n, Set n)] -> Maybe [(n, Set n)]
insert n
a [] = [(n, Set n)] -> Maybe [(n, Set n)]
forall a. a -> Maybe a
Just [(n
a, n -> Set n
deps n
a)]
    insert n
a bs0 :: [(n, Set n)]
bs0@((n
b, Set n
before_b) : [(n, Set n)]
bs)
      | Bool
before Bool -> Bool -> Bool
&& Bool
after = Maybe [(n, Set n)]
forall a. Maybe a
Nothing
      | Bool
before          = ((n
b, Set n -> Set n -> Set n
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set n
before_a Set n
before_b) (n, Set n) -> [(n, Set n)] -> [(n, Set n)]
forall a. a -> [a] -> [a]
:) ([(n, Set n)] -> [(n, Set n)])
-> Maybe [(n, Set n)] -> Maybe [(n, Set n)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> n -> [(n, Set n)] -> Maybe [(n, Set n)]
insert n
a [(n, Set n)]
bs  -- a must come before b
      | Bool
otherwise       = [(n, Set n)] -> Maybe [(n, Set n)]
forall a. a -> Maybe a
Just ([(n, Set n)] -> Maybe [(n, Set n)])
-> [(n, Set n)] -> Maybe [(n, Set n)]
forall a b. (a -> b) -> a -> b
$ (n
a, Set n -> Set n -> Set n
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set n
before_a Set n
before_b) (n, Set n) -> [(n, Set n)] -> [(n, Set n)]
forall a. a -> [a] -> [a]
: [(n, Set n)]
bs0
      where
        before_a :: Set n
before_a = n -> Set n
deps n
a
        before :: Bool
before = n -> Set n -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member n
a Set n
before_b
        after :: Bool
after  = n -> Set n -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member n
b Set n
before_a