{-# LANGUAGE RecordWildCards #-}

module Data.Graph.Sparse.Lowlink where

import Control.Monad
import Control.Monad.ST
import Data.Function
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

import Data.Graph.Sparse
import My.Prelude (rep)

data Lowlink = Lowlink
  { Lowlink -> Vector Int
lowlinkLL :: U.Vector Int
  , Lowlink -> Vector Int
preordLL :: U.Vector Int
  , Lowlink -> Vector Int
parentLL :: U.Vector Vertex
  }
  deriving (Int -> Lowlink -> ShowS
[Lowlink] -> ShowS
Lowlink -> String
(Int -> Lowlink -> ShowS)
-> (Lowlink -> String) -> ([Lowlink] -> ShowS) -> Show Lowlink
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Lowlink -> ShowS
showsPrec :: Int -> Lowlink -> ShowS
$cshow :: Lowlink -> String
show :: Lowlink -> String
$cshowList :: [Lowlink] -> ShowS
showList :: [Lowlink] -> ShowS
Show)

buildLowlink :: SparseGraph w -> Lowlink
buildLowlink :: forall w. SparseGraph w -> Lowlink
buildLowlink SparseGraph w
gr = (forall s. ST s Lowlink) -> Lowlink
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Lowlink) -> Lowlink)
-> (forall s. ST s Lowlink) -> Lowlink
forall a b. (a -> b) -> a -> b
$ do
  let numV :: Int
numV = SparseGraph w -> Int
forall w. SparseGraph w -> Int
numVerticesSG SparseGraph w
gr
  MVector s Int
low <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate Int
numV Int
nothing
  MVector s Int
preord <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate Int
numV Int
nothing
  MVector s Int
parent <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate Int
numV Int
nothing
  MVector s Int
vars <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate Int
numVars Int
0

  Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m ()
rep Int
numV ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
root -> do
    Int
ordRoot <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
preord Int
root
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ordRoot Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nothing) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      ((Int -> Int -> ST s ()) -> Int -> Int -> ST s ())
-> Int -> Int -> ST s ()
forall a. (a -> a) -> a
fix
        ( \Int -> Int -> ST s ()
dfs Int
pv Int
v -> do
            Int
ordV <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
vars Int
_preordId
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
vars Int
_preordId (Int
ordV Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
preord Int
v Int
ordV
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
low Int
v Int
ordV
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
parent Int
v Int
pv

            Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
U.forM_ (SparseGraph w
gr SparseGraph w -> Int -> Vector Int
forall w. SparseGraph w -> Int -> Vector Int
`adj` Int
v) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
u -> do
              Int
ordU <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
preord Int
u
              if Int
ordU Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nothing
                then do
                  Int -> Int -> ST s ()
dfs Int
v Int
u
                  Int
lowU <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
low Int
u
                  MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
UM.unsafeModify MVector s Int
MVector (PrimState (ST s)) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
lowU) Int
v
                else Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
pv) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                  MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
UM.unsafeModify MVector s Int
MVector (PrimState (ST s)) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
ordU) Int
v
        )
        Int
nothing
        Int
root
  Vector Int -> Vector Int -> Vector Int -> Lowlink
Lowlink
    (Vector Int -> Vector Int -> Vector Int -> Lowlink)
-> ST s (Vector Int) -> ST s (Vector Int -> Vector Int -> Lowlink)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
low
    ST s (Vector Int -> Vector Int -> Lowlink)
-> ST s (Vector Int) -> ST s (Vector Int -> Lowlink)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
preord
    ST s (Vector Int -> Lowlink) -> ST s (Vector Int) -> ST s Lowlink
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
parent
  where
    nothing :: Int
nothing = -Int
1
    numVars :: Int
numVars = Int
1
    _preordId :: Int
_preordId = Int
0

articulations :: SparseGraph w -> U.Vector Vertex
articulations :: forall w. SparseGraph w -> Vector Int
articulations SparseGraph w
gr = (Bool -> Bool) -> Vector Bool -> Vector Int
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector Int
U.findIndices Bool -> Bool
forall a. a -> a
id Vector Bool
isArticulation
  where
    nothing :: Int
nothing = -Int
1
    !Lowlink{Vector Int
lowlinkLL :: Lowlink -> Vector Int
preordLL :: Lowlink -> Vector Int
parentLL :: Lowlink -> Vector Int
lowlinkLL :: Vector Int
preordLL :: Vector Int
parentLL :: Vector Int
..} = SparseGraph w -> Lowlink
forall w. SparseGraph w -> Lowlink
buildLowlink SparseGraph w
gr
    !isArticulation :: Vector Bool
isArticulation = (forall s. ST s (MVector s Bool)) -> Vector Bool
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
U.create ((forall s. ST s (MVector s Bool)) -> Vector Bool)
-> (forall s. ST s (MVector s Bool)) -> Vector Bool
forall a b. (a -> b) -> a -> b
$ do
      MVector s Bool
isa <- Int -> Bool -> ST s (MVector (PrimState (ST s)) Bool)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate (SparseGraph w -> Int
forall w. SparseGraph w -> Int
numVerticesSG SparseGraph w
gr) Bool
False
      ((Int -> Int -> ST s ()) -> Vector Int -> ST s ())
-> Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> Int -> ST s ()) -> Vector Int -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Int -> a -> m b) -> Vector a -> m ()
U.imapM_ Vector Int
parentLL ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v Int
pv -> Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
pv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nothing) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
preordLL Int
pv Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
lowlinkLL Int
v) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector s Bool
MVector (PrimState (ST s)) Bool
isa Int
pv Bool
True
      Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
U.forM_ ((Int -> Bool) -> Vector Int -> Vector Int
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector Int
U.findIndices (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nothing) Vector Int
parentLL) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
root -> do
        MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector s Bool
MVector (PrimState (ST s)) Bool
isa Int
root
          (Bool -> ST s ()) -> (Vector Int -> Bool) -> Vector Int -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2)
          (Int -> Bool) -> (Vector Int -> Int) -> Vector Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int
forall a. Unbox a => Vector a -> Int
U.length
          (Vector Int -> Int)
-> (Vector Int -> Vector Int) -> Vector Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Bool) -> Vector Int -> Vector Int
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
U.filter (\Int
v -> Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
parentLL Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
root)
          (Vector Int -> ST s ()) -> Vector Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ SparseGraph w
gr SparseGraph w -> Int -> Vector Int
forall w. SparseGraph w -> Int -> Vector Int
`adj` Int
root
      MVector s Bool -> ST s (MVector s Bool)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Bool
isa

bridges :: SparseGraph w -> U.Vector Edge
bridges :: forall w. SparseGraph w -> Vector Edge
bridges SparseGraph w
gr = (Int -> Int -> Maybe Edge) -> Vector Int -> Vector Edge
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> Maybe b) -> Vector a -> Vector b
U.imapMaybe Int -> Int -> Maybe Edge
isBridge Vector Int
parentLL
  where
    nothing :: Int
nothing = -Int
1
    !Lowlink{Vector Int
lowlinkLL :: Lowlink -> Vector Int
preordLL :: Lowlink -> Vector Int
parentLL :: Lowlink -> Vector Int
parentLL :: Vector Int
lowlinkLL :: Vector Int
preordLL :: Vector Int
..} = SparseGraph w -> Lowlink
forall w. SparseGraph w -> Lowlink
buildLowlink SparseGraph w
gr
    isBridge :: Int -> Int -> Maybe Edge
isBridge Int
v Int
pv
      | Int
pv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nothing
      , Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
preordLL Int
pv Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
lowlinkLL Int
v =
          Edge -> Maybe Edge
forall a. a -> Maybe a
Just (Int
pv, Int
v)
      | Bool
otherwise = Maybe Edge
forall a. Maybe a
Nothing