{-# 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 qualified Data.Graph as Graph
import Control.Arrow
import Agda.Utils.List (nubOn)
import Agda.Utils.SemiRing
import qualified Agda.Utils.Graph.AdjacencyMap.Unidirectional as G

mergeBy :: (a -> a -> Bool) -> [a] -> [a] -> [a]
mergeBy :: (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 :: [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