{-# LANGUAGE RecordWildCards #-}
{-# OPTIONS_HADDOCK hide #-}

-- | Implementation of Strongly Connected Components calculation. Use `AtCoder.Scc` instead.
--
-- @since 1.0.0.0
module AtCoder.Internal.Scc
  ( -- * Internal SCC
    SccGraph (nScc),

    -- * Constructors
    new,

    -- * Adding edges
    addEdge,

    -- * SCC calculation
    sccIds,
    scc,

    -- ** (Extra API) CSR API
    sccIdsCsr,
    sccCsr
  )
where

import AtCoder.Internal.Csr qualified as ACICSR
import AtCoder.Internal.GrowVec qualified as ACIGV
import Control.Monad (unless, when)
import Control.Monad.Fix (fix)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST (runST)
import Data.Foldable (for_)
import Data.Maybe (fromJust)
import Data.Vector qualified as V
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM

-- | Graph for collecting strongly connected components.
--
-- @since 1.0.0.0
data SccGraph s = SccGraph
  { -- | The number of vertices.
    --
    -- @since 1.0.0.0
    forall s. SccGraph s -> Int
nScc :: {-# UNPACK #-} !Int,
    forall s. SccGraph s -> GrowVec s (Int, Int)
edgesScc :: !(ACIGV.GrowVec s (Int, Int))
  }

-- | \(O(n)\) Creates a `SccGraph` of \(n\) vertices.
--
-- @since 1.0.0.0
{-# INLINE new #-}
new :: (PrimMonad m) => Int -> m (SccGraph (PrimState m))
new :: forall (m :: * -> *).
PrimMonad m =>
Int -> m (SccGraph (PrimState m))
new Int
nScc = do
  GrowVec (PrimState m) (Int, Int)
edgesScc <- Int -> m (GrowVec (PrimState m) (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
ACIGV.new Int
0
  SccGraph (PrimState m) -> m (SccGraph (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SccGraph {Int
GrowVec (PrimState m) (Int, Int)
nScc :: Int
edgesScc :: GrowVec (PrimState m) (Int, Int)
nScc :: Int
edgesScc :: GrowVec (PrimState m) (Int, Int)
..}

-- | \(O(1)\) amortized. Adds an edge to the graph.
--
-- @since 1.0.0.0
{-# INLINE addEdge #-}
addEdge :: (PrimMonad m) => SccGraph (PrimState m) -> Int -> Int -> m ()
addEdge :: forall (m :: * -> *).
PrimMonad m =>
SccGraph (PrimState m) -> Int -> Int -> m ()
addEdge SccGraph {GrowVec (PrimState m) (Int, Int)
edgesScc :: forall s. SccGraph s -> GrowVec s (Int, Int)
edgesScc :: GrowVec (PrimState m) (Int, Int)
edgesScc} Int
from Int
to = do
  GrowVec (PrimState m) (Int, Int) -> (Int, Int) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
ACIGV.pushBack GrowVec (PrimState m) (Int, Int)
edgesScc (Int
from, Int
to)

-- | \(O(n + m)\) Returns a pair of @(# of scc, scc id)@.
--
-- @since 1.0.0.0
{-# INLINE sccIds #-}
sccIds :: (PrimMonad m) => SccGraph (PrimState m) -> m (Int, VU.Vector Int)
sccIds :: forall (m :: * -> *).
PrimMonad m =>
SccGraph (PrimState m) -> m (Int, Vector Int)
sccIds SccGraph {Int
GrowVec (PrimState m) (Int, Int)
nScc :: forall s. SccGraph s -> Int
edgesScc :: forall s. SccGraph s -> GrowVec s (Int, Int)
nScc :: Int
edgesScc :: GrowVec (PrimState m) (Int, Int)
..} = do
  Csr ()
csr <- HasCallStack => Int -> Vector (Int, Int) -> Csr ()
Int -> Vector (Int, Int) -> Csr ()
ACICSR.build' Int
nScc (Vector (Int, Int) -> Csr ())
-> m (Vector (Int, Int)) -> m (Csr ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GrowVec (PrimState m) (Int, Int) -> m (Vector (Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Vector a)
ACIGV.unsafeFreeze GrowVec (PrimState m) (Int, Int)
edgesScc
  (Int, Vector Int) -> m (Int, Vector Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, Vector Int) -> m (Int, Vector Int))
-> (Int, Vector Int) -> m (Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Csr () -> (Int, Vector Int)
forall w. Csr w -> (Int, Vector Int)
sccIdsCsr Csr ()
csr

-- | \(O(n + m)\) Returns the strongly connected components.
--
-- @since 1.0.0.0
{-# INLINE scc #-}
scc :: (PrimMonad m) => SccGraph (PrimState m) -> m (V.Vector (VU.Vector Int))
scc :: forall (m :: * -> *).
PrimMonad m =>
SccGraph (PrimState m) -> m (Vector (Vector Int))
scc SccGraph (PrimState m)
g = do
  (!Int
groupNum, !Vector Int
ids) <- SccGraph (PrimState m) -> m (Int, Vector Int)
forall (m :: * -> *).
PrimMonad m =>
SccGraph (PrimState m) -> m (Int, Vector Int)
sccIds SccGraph (PrimState m)
g
  let counts :: Vector Int
counts = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        MVector s Int
vec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
groupNum (Int
0 :: Int)
        Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
ids ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
          MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
vec (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
x
        MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
vec
  Vector (MVector (PrimState m) Int)
groups <- (Int -> m (MVector (PrimState m) Int))
-> Vector Int -> m (Vector (MVector (PrimState m) Int))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Vector Int -> m (Vector (MVector (PrimState m) Int)))
-> Vector Int -> m (Vector (MVector (PrimState m) Int))
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert Vector Int
counts
  MVector (PrimState m) Int
is <- Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
groupNum (Int
0 :: Int)
  Vector Int -> (Int -> Int -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
ids ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
v Int
sccId -> do
    Int
i <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Int
is Int
sccId
    MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Int
is Int
sccId (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write (Vector (MVector (PrimState m) Int)
groups Vector (MVector (PrimState m) Int)
-> Int -> MVector (PrimState m) Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
sccId) Int
i Int
v
  (MVector (PrimState m) Int -> m (Vector Int))
-> Vector (MVector (PrimState m) Int) -> m (Vector (Vector Int))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM MVector (PrimState m) Int -> m (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze Vector (MVector (PrimState m) Int)
groups

-- | \(O(n + m)\) API) Returns a pair of @(# of scc, scc id)@.
--
-- @since 1.1.0.0
{-# INLINE sccIdsCsr #-}
sccIdsCsr :: ACICSR.Csr w -> (Int, VU.Vector Int)
sccIdsCsr :: forall w. Csr w -> (Int, Vector Int)
sccIdsCsr g :: Csr w
g@ACICSR.Csr {Int
Vector w
Vector Int
nCsr :: Int
mCsr :: Int
startCsr :: Vector Int
adjCsr :: Vector Int
wCsr :: Vector w
nCsr :: forall w. Csr w -> Int
mCsr :: forall w. Csr w -> Int
startCsr :: forall w. Csr w -> Vector Int
adjCsr :: forall w. Csr w -> Vector Int
wCsr :: forall w. Csr w -> Vector w
..} = (forall s. ST s (Int, Vector Int)) -> (Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Int, Vector Int)) -> (Int, Vector Int))
-> (forall s. ST s (Int, Vector Int)) -> (Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  -- see also the Wikipedia: https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm#The_algorithm_in_pseudocode
  -- next SCC ID
  MVector s Int
groupNum <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (Int
0 :: Int)
  -- stack of vertices
  GrowVec s Int
visited <- Int -> ST s (GrowVec (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (GrowVec (PrimState m) a)
ACIGV.new Int
nCsr
  -- vertex -> low-link: the smallest index of any node on the stack known to be reachable from
  -- v through v's DFS subtree, including v itself.
  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)
VUM.replicate Int
nCsr (Int
0 :: Int)
  -- vertex -> order of the visit (0, 1, ..)
  MVector s Int
ord <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nCsr (-Int
1 :: Int)
  -- vertex -> scc id
  MVector s Int
ids <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nCsr (Int
0 :: Int)

  let dfs :: Int -> Int -> ST s Int
dfs Int
v Int
ord0 = do
        MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
low Int
v Int
ord0
        MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
ord Int
v Int
ord0
        GrowVec (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> a -> m ()
ACIGV.pushBack GrowVec s Int
GrowVec (PrimState (ST s)) Int
visited Int
v
        -- look around @v@, folding their low-link onto the low-link of @v@.
        Int
ord' <-
          (Int -> Int -> ST s Int) -> Int -> Vector Int -> ST s Int
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
            ( \Int
curOrd Int
to -> do
                Int
ordTo <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
ord Int
to
                if Int
ordTo Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1
                  then do
                    -- not visited yet.
                    Int
nextOrd <- Int -> Int -> ST s Int
dfs Int
to (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
curOrd
                    Int
lowTo <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
low Int
to
                    MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
lowTo) Int
v
                    Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
nextOrd
                  else do
                    -- lookup back and update the low-link.
                    MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
ordTo) Int
v
                    Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
curOrd
            )
            (Int
ord0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            (Csr w
g Csr w -> Int -> Vector Int
forall w. HasCallStack => Csr w -> Int -> Vector Int
`ACICSR.adj` Int
v)

        Int
lowV <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
low Int
v
        Int
ordV <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
ord Int
v
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lowV Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ordV) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          -- it's the root of a SCC, no more to look back
          Int
sccId <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
groupNum Int
0
          (ST s () -> ST s ()) -> ST s ()
forall a. (a -> a) -> a
fix ((ST s () -> ST s ()) -> ST s ())
-> (ST s () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ST s ()
loop -> do
            Int
u <- Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> ST s (Maybe Int) -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GrowVec (PrimState (ST s)) Int -> ST s (Maybe Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
GrowVec (PrimState m) a -> m (Maybe a)
ACIGV.popBack GrowVec s Int
GrowVec (PrimState (ST s)) Int
visited
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
ord Int
u Int
nCsr
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
ids Int
u Int
sccId
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
v) ST s ()
loop
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
groupNum Int
0 (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
sccId Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ord'

  (Int -> Int -> ST s Int) -> Int -> Vector Int -> ST s ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m ()
VU.foldM'_
    ( \Int
curOrd Int
i -> do
        Int
o <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
ord Int
i
        if Int
o Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1
          then Int -> Int -> ST s Int
dfs Int
i Int
curOrd
          else Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
curOrd
    )
    (Int
0 :: Int)
    (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
nCsr Int -> Int
forall a. a -> a
id)

  Int
num <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
groupNum Int
0
  -- The SCCs are reverse topologically sorted, e.g., [0, 1] <- [2] <- [3]
  -- Now reverse the SCC IDs so that they will be topologically sorted: [3] -> [2] -> [0, 1]
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
nCsr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
ids ((Int
num Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) -) Int
i

  Vector Int
ids' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
ids
  (Int, Vector Int) -> ST s (Int, Vector Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
num, Vector Int
ids')

-- | \(O(n + m)\) Returns the strongly connected components.
--
-- @since 1.1.0.0
{-# INLINE sccCsr #-}
sccCsr :: ACICSR.Csr w -> V.Vector (VU.Vector Int)
sccCsr :: forall w. Csr w -> Vector (Vector Int)
sccCsr Csr w
g = (forall s. ST s (Vector (Vector Int))) -> Vector (Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector (Vector Int))) -> Vector (Vector Int))
-> (forall s. ST s (Vector (Vector Int))) -> Vector (Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  Vector (MVector s Int)
groups <- (Int -> ST s (MVector s Int))
-> Vector Int -> ST s (Vector (MVector s Int))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM Int -> ST s (MVector s Int)
Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Vector Int -> ST s (Vector (MVector s Int)))
-> Vector Int -> ST s (Vector (MVector s Int))
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert Vector Int
counts
  MVector s Int
is <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
groupNum (Int
0 :: Int)
  Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
ids ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v Int
sccId -> do
    Int
i <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
is Int
sccId
    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
is Int
sccId (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write (Vector (MVector s Int)
groups Vector (MVector s Int) -> Int -> MVector s Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
sccId) Int
i Int
v
  (MVector s Int -> ST s (Vector Int))
-> Vector (MVector s Int) -> ST s (Vector (Vector Int))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM MVector s Int -> ST s (Vector Int)
MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze Vector (MVector s Int)
groups
  where
    (!Int
groupNum, !Vector Int
ids) = Csr w -> (Int, Vector Int)
forall w. Csr w -> (Int, Vector Int)
sccIdsCsr Csr w
g
    counts :: Vector Int
counts = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
      MVector s Int
vec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
groupNum (Int
0 :: Int)
      Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
ids ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
        MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
vec (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
x
      MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
vec