{-# LANGUAGE LambdaCase #-}
module Algebra.Graph.AdjacencyIntMap.Algorithm (
bfsForest, bfs, dfsForest, dfsForestFrom, dfs, reachable,
topSort, isAcyclic,
isDfsForestOf, isTopSortOf,
Cycle
) where
import Control.Monad
import Control.Monad.Trans.Cont
import Control.Monad.Trans.State.Strict
import Data.Either
import Data.List.NonEmpty (NonEmpty(..), (<|))
import Data.Tree
import Algebra.Graph.AdjacencyIntMap
import qualified Data.List as List
import qualified Data.IntMap.Strict as IntMap
import qualified Data.IntSet as IntSet
bfsForest :: AdjacencyIntMap -> [Int] -> Forest Int
bfsForest :: AdjacencyIntMap -> [Int] -> Forest Int
bfsForest AdjacencyIntMap
g [Int]
vs= forall s a. State s a -> s -> a
evalState ([Int] -> StateT IntSet Identity (Forest Int)
explore [ Int
v | Int
v <- [Int]
vs, Int -> AdjacencyIntMap -> Bool
hasVertex Int
v AdjacencyIntMap
g ]) IntSet
IntSet.empty
where
explore :: [Int] -> StateT IntSet Identity (Forest Int)
explore = forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM forall {m :: * -> *}. Monad m => Int -> StateT IntSet m Bool
discovered forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (m :: * -> *) b a.
Monad m =>
(b -> m (a, [b])) -> [b] -> m [Tree a]
unfoldForestM_BF Int -> StateT IntSet Identity (Int, [Int])
walk
walk :: Int -> StateT IntSet Identity (Int, [Int])
walk Int
v = (Int
v,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> StateT IntSet Identity [Int]
adjacentM Int
v
adjacentM :: Int -> StateT IntSet Identity [Int]
adjacentM Int
v = forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM forall {m :: * -> *}. Monad m => Int -> StateT IntSet m Bool
discovered forall a b. (a -> b) -> a -> b
$ IntSet -> [Int]
IntSet.toList (Int -> AdjacencyIntMap -> IntSet
postIntSet Int
v AdjacencyIntMap
g)
discovered :: Int -> StateT IntSet m Bool
discovered Int
v = do Bool
new <- forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> IntSet -> Bool
IntSet.member Int
v)
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
new forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' (Int -> IntSet -> IntSet
IntSet.insert Int
v)
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
new
bfs :: AdjacencyIntMap -> [Int] -> [[Int]]
bfs :: AdjacencyIntMap -> [Int] -> [[Int]]
bfs AdjacencyIntMap
g = forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [[a]] -> [[a]]
List.transpose forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Tree a -> [[a]]
levels forall b c a. (b -> c) -> (a -> b) -> a -> c
. AdjacencyIntMap -> [Int] -> Forest Int
bfsForest AdjacencyIntMap
g
dfsForestFromImpl :: AdjacencyIntMap -> [Int] -> Forest Int
dfsForestFromImpl :: AdjacencyIntMap -> [Int] -> Forest Int
dfsForestFromImpl AdjacencyIntMap
g [Int]
vs = forall s a. State s a -> s -> a
evalState ([Int] -> StateT IntSet Identity (Forest Int)
explore [Int]
vs) IntSet
IntSet.empty
where
explore :: [Int] -> StateT IntSet Identity (Forest Int)
explore (Int
v:[Int]
vs) = forall {m :: * -> *}. Monad m => Int -> StateT IntSet m Bool
discovered Int
v forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Bool
True -> (:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> StateT IntSet Identity (Tree Int)
walk Int
v forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Int] -> StateT IntSet Identity (Forest Int)
explore [Int]
vs
Bool
False -> [Int] -> StateT IntSet Identity (Forest Int)
explore [Int]
vs
explore [] = forall (m :: * -> *) a. Monad m => a -> m a
return []
walk :: Int -> StateT IntSet Identity (Tree Int)
walk Int
v = forall a. a -> [Tree a] -> Tree a
Node Int
v forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> StateT IntSet Identity (Forest Int)
explore (Int -> [Int]
adjacent Int
v)
adjacent :: Int -> [Int]
adjacent Int
v = IntSet -> [Int]
IntSet.toList (Int -> AdjacencyIntMap -> IntSet
postIntSet Int
v AdjacencyIntMap
g)
discovered :: Int -> StateT IntSet m Bool
discovered Int
v = do Bool
new <- forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> IntSet -> Bool
IntSet.member Int
v)
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
new forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' (Int -> IntSet -> IntSet
IntSet.insert Int
v)
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
new
dfsForest :: AdjacencyIntMap -> Forest Int
dfsForest :: AdjacencyIntMap -> Forest Int
dfsForest AdjacencyIntMap
g = AdjacencyIntMap -> [Int] -> Forest Int
dfsForestFromImpl AdjacencyIntMap
g (AdjacencyIntMap -> [Int]
vertexList AdjacencyIntMap
g)
dfsForestFrom :: AdjacencyIntMap -> [Int] -> Forest Int
dfsForestFrom :: AdjacencyIntMap -> [Int] -> Forest Int
dfsForestFrom AdjacencyIntMap
g [Int]
vs = AdjacencyIntMap -> [Int] -> Forest Int
dfsForestFromImpl AdjacencyIntMap
g [ Int
v | Int
v <- [Int]
vs, Int -> AdjacencyIntMap -> Bool
hasVertex Int
v AdjacencyIntMap
g ]
dfs :: AdjacencyIntMap -> [Int] -> [Int]
dfs :: AdjacencyIntMap -> [Int] -> [Int]
dfs AdjacencyIntMap
x = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Tree a -> [a]
flatten forall b c a. (b -> c) -> (a -> b) -> a -> c
. AdjacencyIntMap -> [Int] -> Forest Int
dfsForestFrom AdjacencyIntMap
x
reachable :: AdjacencyIntMap -> Int -> [Int]
reachable :: AdjacencyIntMap -> Int -> [Int]
reachable AdjacencyIntMap
x Int
y = AdjacencyIntMap -> [Int] -> [Int]
dfs AdjacencyIntMap
x [Int
y]
type Cycle = NonEmpty
type Result = Either (Cycle Int) [Int]
data NodeState = Entered | Exited
data S = S { S -> IntMap Int
parent :: IntMap.IntMap Int
, S -> IntMap NodeState
entry :: IntMap.IntMap NodeState
, S -> [Int]
order :: [Int] }
topSortImpl :: AdjacencyIntMap -> StateT S (Cont Result) Result
topSortImpl :: AdjacencyIntMap -> StateT S (Cont Result) Result
topSortImpl AdjacencyIntMap
g = forall (m :: * -> *) a s b.
CallCC m (a, s) (b, s) -> CallCC (StateT s m) a b
liftCallCC' forall {k} a (r :: k) (m :: k -> *) b.
((a -> ContT r m b) -> ContT r m a) -> ContT r m a
callCC forall a b. (a -> b) -> a -> b
$ \Result -> StateT S (Cont Result) ()
cyclic ->
do let vertices :: [Int]
vertices = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. IntMap a -> [(Int, a)]
IntMap.toDescList forall a b. (a -> b) -> a -> b
$ AdjacencyIntMap -> IntMap IntSet
adjacencyIntMap AdjacencyIntMap
g
adjacent :: Int -> [Int]
adjacent = IntSet -> [Int]
IntSet.toDescList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> AdjacencyIntMap -> IntSet
postIntSet AdjacencyIntMap
g
dfsRoot :: Int -> StateT S (Cont Result) ()
dfsRoot Int
x = forall {m :: * -> *}.
Monad m =>
Int -> StateT S m (Maybe NodeState)
nodeState Int
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Maybe NodeState
Nothing -> forall {m :: * -> *}. Monad m => Int -> StateT S m ()
enterRoot Int
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> StateT S (Cont Result) ()
dfs Int
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {m :: * -> *}. Monad m => Int -> StateT S m ()
exit Int
x
Maybe NodeState
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
dfs :: Int -> StateT S (Cont Result) ()
dfs Int
x = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Int -> [Int]
adjacent Int
x) forall a b. (a -> b) -> a -> b
$ \Int
y ->
forall {m :: * -> *}.
Monad m =>
Int -> StateT S m (Maybe NodeState)
nodeState Int
y forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Maybe NodeState
Nothing -> forall {m :: * -> *}. Monad m => Int -> Int -> StateT S m ()
enter Int
x Int
y forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> StateT S (Cont Result) ()
dfs Int
y forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {m :: * -> *}. Monad m => Int -> StateT S m ()
exit Int
y
Just NodeState
Exited -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
Just NodeState
Entered -> Result -> StateT S (Cont Result) ()
cyclic forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> IntMap Int -> NonEmpty Int
retrace Int
x Int
y forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets S -> IntMap Int
parent
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int]
vertices Int -> StateT S (Cont Result) ()
dfsRoot
forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets S -> [Int]
order
where
nodeState :: Int -> StateT S m (Maybe NodeState)
nodeState Int
v = forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets (forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. S -> IntMap NodeState
entry)
enter :: Int -> Int -> StateT S m ()
enter Int
u Int
v = forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' (\(S IntMap Int
m IntMap NodeState
n [Int]
vs) -> IntMap Int -> IntMap NodeState -> [Int] -> S
S (forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
v Int
u IntMap Int
m)
(forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
v NodeState
Entered IntMap NodeState
n)
[Int]
vs)
enterRoot :: Int -> StateT S m ()
enterRoot Int
v = forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' (\(S IntMap Int
m IntMap NodeState
n [Int]
vs) -> IntMap Int -> IntMap NodeState -> [Int] -> S
S IntMap Int
m (forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
v NodeState
Entered IntMap NodeState
n) [Int]
vs)
exit :: Int -> StateT S m ()
exit Int
v = forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify' (\(S IntMap Int
m IntMap NodeState
n [Int]
vs) -> IntMap Int -> IntMap NodeState -> [Int] -> S
S IntMap Int
m (forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IntMap.alter (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NodeState -> NodeState
leave) Int
v IntMap NodeState
n) (Int
vforall a. a -> [a] -> [a]
:[Int]
vs))
where leave :: NodeState -> NodeState
leave = \case
NodeState
Entered -> NodeState
Exited
NodeState
Exited -> forall a. HasCallStack => [Char] -> a
error [Char]
"Internal error: dfs search order violated"
retrace :: Int -> Int -> IntMap Int -> NonEmpty Int
retrace Int
curr Int
head IntMap Int
parent = NonEmpty Int -> NonEmpty Int
aux (Int
curr forall a. a -> [a] -> NonEmpty a
:| []) where
aux :: NonEmpty Int -> NonEmpty Int
aux xs :: NonEmpty Int
xs@(Int
curr :| [Int]
_)
| Int
head forall a. Eq a => a -> a -> Bool
== Int
curr = NonEmpty Int
xs
| Bool
otherwise = NonEmpty Int -> NonEmpty Int
aux (IntMap Int
parent forall a. IntMap a -> Int -> a
IntMap.! Int
curr forall a. a -> NonEmpty a -> NonEmpty a
<| NonEmpty Int
xs)
topSort :: AdjacencyIntMap -> Either (Cycle Int) [Int]
topSort :: AdjacencyIntMap -> Result
topSort AdjacencyIntMap
g = forall r a. Cont r a -> (a -> r) -> r
runCont (forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (AdjacencyIntMap -> StateT S (Cont Result) Result
topSortImpl AdjacencyIntMap
g) S
initialState) forall a. a -> a
id
where
initialState :: S
initialState = IntMap Int -> IntMap NodeState -> [Int] -> S
S forall a. IntMap a
IntMap.empty forall a. IntMap a
IntMap.empty []
isAcyclic :: AdjacencyIntMap -> Bool
isAcyclic :: AdjacencyIntMap -> Bool
isAcyclic = forall a b. Either a b -> Bool
isRight forall b c a. (b -> c) -> (a -> b) -> a -> c
. AdjacencyIntMap -> Result
topSort
isDfsForestOf :: Forest Int -> AdjacencyIntMap -> Bool
isDfsForestOf :: Forest Int -> AdjacencyIntMap -> Bool
isDfsForestOf Forest Int
f AdjacencyIntMap
am = case IntSet -> Forest Int -> Maybe IntSet
go IntSet
IntSet.empty Forest Int
f of
Just IntSet
seen -> IntSet
seen forall a. Eq a => a -> a -> Bool
== AdjacencyIntMap -> IntSet
vertexIntSet AdjacencyIntMap
am
Maybe IntSet
Nothing -> Bool
False
where
go :: IntSet -> Forest Int -> Maybe IntSet
go IntSet
seen [] = forall a. a -> Maybe a
Just IntSet
seen
go IntSet
seen (Tree Int
t:Forest Int
ts) = do
let root :: Int
root = forall a. Tree a -> a
rootLabel Tree Int
t
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ Int
root Int -> IntSet -> Bool
`IntSet.notMember` IntSet
seen
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Int -> Int -> AdjacencyIntMap -> Bool
hasEdge Int
root (forall a. Tree a -> a
rootLabel Tree Int
subTree) AdjacencyIntMap
am | Tree Int
subTree <- forall a. Tree a -> [Tree a]
subForest Tree Int
t ]
IntSet
newSeen <- IntSet -> Forest Int -> Maybe IntSet
go (Int -> IntSet -> IntSet
IntSet.insert Int
root IntSet
seen) (forall a. Tree a -> [Tree a]
subForest Tree Int
t)
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ Int -> AdjacencyIntMap -> IntSet
postIntSet Int
root AdjacencyIntMap
am IntSet -> IntSet -> Bool
`IntSet.isSubsetOf` IntSet
newSeen
IntSet -> Forest Int -> Maybe IntSet
go IntSet
newSeen Forest Int
ts
isTopSortOf :: [Int] -> AdjacencyIntMap -> Bool
isTopSortOf :: [Int] -> AdjacencyIntMap -> Bool
isTopSortOf [Int]
xs AdjacencyIntMap
m = IntSet -> [Int] -> Bool
go IntSet
IntSet.empty [Int]
xs
where
go :: IntSet -> [Int] -> Bool
go IntSet
seen [] = IntSet
seen forall a. Eq a => a -> a -> Bool
== forall a. IntMap a -> IntSet
IntMap.keysSet (AdjacencyIntMap -> IntMap IntSet
adjacencyIntMap AdjacencyIntMap
m)
go IntSet
seen (Int
v:[Int]
vs) = Int -> AdjacencyIntMap -> IntSet
postIntSet Int
v AdjacencyIntMap
m IntSet -> IntSet -> IntSet
`IntSet.intersection` IntSet
newSeen forall a. Eq a => a -> a -> Bool
== IntSet
IntSet.empty
Bool -> Bool -> Bool
&& IntSet -> [Int] -> Bool
go IntSet
newSeen [Int]
vs
where
newSeen :: IntSet
newSeen = Int -> IntSet -> IntSet
IntSet.insert Int
v IntSet
seen