{-# 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
  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
  preord <- UM.replicate numV nothing
  parent <- UM.replicate numV nothing
  vars <- UM.replicate numVars 0

  rep numV $ \Int
root -> do
    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
    when (ordRoot == nothing) $ do
      fix
        ( \Int -> Int -> ST s ()
dfs Int
pv Int
v -> do
            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
            UM.unsafeWrite vars _preordId (ordV + 1)

            UM.unsafeWrite preord v ordV
            UM.unsafeWrite low v ordV
            UM.unsafeWrite parent v pv

            U.forM_ (gr `adj` v) $ \Int
u -> do
              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 ordU == nothing
                then do
                  dfs v u
                  lowU <- UM.unsafeRead low u
                  UM.unsafeModify low (min lowU) v
                else when (u /= pv) $ do
                  UM.unsafeModify low (min ordU) v
        )
        nothing
        root
  Lowlink
    <$> U.unsafeFreeze low
    <*> U.unsafeFreeze preord
    <*> U.unsafeFreeze 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
      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
      flip U.imapM_ parentLL $ \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
      U.forM_ (U.findIndices (== nothing) parentLL) $ \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
      return 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