{-# LANGUAGE Rank2Types, Trustworthy #-}
module Data.Graph.ArraySCC(scc) where
import Data.Graph(Graph,Vertex)
import Data.Array.ST(STUArray, newArray, readArray, writeArray)
import Data.Array as A
import Data.Array.Unsafe(unsafeFreeze)
import Control.Monad.ST
import Control.Monad(ap)
scc :: Graph -> ([(Int,[Vertex])], Vertex -> Int)
scc :: Graph -> ([(Int, [Int])], Int -> Int)
scc g :: Graph
g = (forall s. ST s ([(Int, [Int])], Int -> Int))
-> ([(Int, [Int])], Int -> Int)
forall a. (forall s. ST s a) -> a
runST (
do STUArray s Int Int
ixes <- (Int, Int) -> Int -> ST s (STUArray s Int Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Graph -> (Int, Int)
forall i e. Array i e -> (i, i)
bounds Graph
g) 0
STUArray s Int Int
lows <- (Int, Int) -> Int -> ST s (STUArray s Int Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Graph -> (Int, Int)
forall i e. Array i e -> (i, i)
bounds Graph
g) 0
S
s <- Func s ([Int] -> ST s S)
forall s. Func s ([Int] -> ST s S)
roots Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows ([Int] -> Int -> [(Int, [Int])] -> Int -> S
S [] 1 [] 1) (Graph -> [Int]
forall i e. Ix i => Array i e -> [i]
indices Graph
g)
Array Int Int
sccm <- STUArray s Int Int -> ST s (Array Int Int)
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze STUArray s Int Int
ixes
([(Int, [Int])], Int -> Int) -> ST s ([(Int, [Int])], Int -> Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (S -> [(Int, [Int])]
sccs S
s, \i :: Int
i -> Array Int Int
sccm Array Int Int -> Int -> Int
forall i e. Ix i => Array i e -> i -> e
! Int
i)
)
type Func s a =
Graph
-> STUArray s Vertex Int
-> STUArray s Vertex Int
-> S
-> a
data S = S { S -> [Int]
stack :: ![Vertex]
, S -> Int
num :: !Int
, S -> [(Int, [Int])]
sccs :: ![(Int,[Vertex])]
, S -> Int
next_scc :: !Int
}
roots :: Func s ([Vertex] -> ST s S)
roots :: Func s ([Int] -> ST s S)
roots g :: Graph
g ixes :: STUArray s Int Int
ixes lows :: STUArray s Int Int
lows st :: S
st (v :: Int
v:vs :: [Int]
vs) =
do Int
i <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
ixes Int
v
if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 then do S
s1 <- Func s (Int -> ST s S)
forall s. Func s (Int -> ST s S)
from_root Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows S
st Int
v
Func s ([Int] -> ST s S)
forall s. Func s ([Int] -> ST s S)
roots Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows S
s1 [Int]
vs
else Func s ([Int] -> ST s S)
forall s. Func s ([Int] -> ST s S)
roots Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows S
st [Int]
vs
roots _ _ _ s :: S
s [] = S -> ST s S
forall (m :: * -> *) a. Monad m => a -> m a
return S
s
from_root :: Func s (Vertex -> ST s S)
from_root :: Func s (Int -> ST s S)
from_root g :: Graph
g ixes :: STUArray s Int Int
ixes lows :: STUArray s Int Int
lows s :: S
s v :: Int
v =
do let me :: Int
me = S -> Int
num S
s
STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
ixes Int
v (Int -> Int
forall a. Num a => a -> a
negate Int
me)
STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
lows Int
v Int
me
S
newS <- Func s (Int -> [Int] -> ST s S)
forall s. Func s (Int -> [Int] -> ST s S)
check_adj Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows
S
s { stack :: [Int]
stack = Int
v Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: S -> [Int]
stack S
s, num :: Int
num = Int
me Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1 } Int
v (Graph
g Graph -> Int -> [Int]
forall i e. Ix i => Array i e -> i -> e
! Int
v)
Int
x <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
lows Int
v
if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
me then S -> ST s S
forall (m :: * -> *) a. Monad m => a -> m a
return S
newS else
case (Int -> Bool) -> [Int] -> ([Int], [Int])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
v) (S -> [Int]
stack S
newS) of
(as :: [Int]
as,b :: Int
b:bs :: [Int]
bs) ->
do let this :: [Int]
this = Int
b Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
as
n :: Int
n = S -> Int
next_scc S
newS
(Int -> ST s ()) -> [Int] -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\i :: Int
i -> STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
ixes Int
i Int
n) [Int]
this
S -> ST s S
forall (m :: * -> *) a. Monad m => a -> m a
return $WS :: [Int] -> Int -> [(Int, [Int])] -> Int -> S
S { stack :: [Int]
stack = [Int]
bs
, num :: Int
num = S -> Int
num S
newS
, sccs :: [(Int, [Int])]
sccs = (Int
n,[Int]
this) (Int, [Int]) -> [(Int, [Int])] -> [(Int, [Int])]
forall a. a -> [a] -> [a]
: S -> [(Int, [Int])]
sccs S
newS
, next_scc :: Int
next_scc = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1
}
_ -> [Char] -> ST s S
forall a. HasCallStack => [Char] -> a
error ("bug in scc---vertex not on the stack: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
v)
check_adj :: Func s (Vertex -> [Vertex] -> ST s S)
check_adj :: Func s (Int -> [Int] -> ST s S)
check_adj g :: Graph
g ixes :: STUArray s Int Int
ixes lows :: STUArray s Int Int
lows st :: S
st v :: Int
v (v' :: Int
v':vs :: [Int]
vs) =
do Int
i <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
ixes Int
v'
case () of
_ | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 ->
do S
newS <- Func s (Int -> ST s S)
forall s. Func s (Int -> ST s S)
from_root Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows S
st Int
v'
Int
new_low <- Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int) -> ST s Int -> ST s (Int -> Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
lows Int
v ST s (Int -> Int) -> ST s Int -> ST s Int
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
`ap` STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
lows Int
v'
STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
lows Int
v Int
new_low
Func s (Int -> [Int] -> ST s S)
forall s. Func s (Int -> [Int] -> ST s S)
check_adj Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows S
newS Int
v [Int]
vs
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< 0 ->
do Int
j <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
lows Int
v
STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
lows Int
v (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
j (Int -> Int
forall a. Num a => a -> a
negate Int
i))
Func s (Int -> [Int] -> ST s S)
forall s. Func s (Int -> [Int] -> ST s S)
check_adj Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows S
st Int
v [Int]
vs
| Bool
otherwise -> Func s (Int -> [Int] -> ST s S)
forall s. Func s (Int -> [Int] -> ST s S)
check_adj Graph
g STUArray s Int Int
ixes STUArray s Int Int
lows S
st Int
v [Int]
vs
check_adj _ _ _ st :: S
st _ [] = S -> ST s S
forall (m :: * -> *) a. Monad m => a -> m a
return S
st