{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

module Data.Graph.Sparse.BFS where

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

import Data.Buffer
import Data.Graph.Sparse

-- | /O(V+E)/
bfsSG :: Vertex -> SparseGraph w -> U.Vector Int
bfsSG :: forall w. Vertex -> SparseGraph w -> Vector Vertex
bfsSG Vertex
source gr :: SparseGraph w
gr@SparseGraph{Vertex
Vector w
Vector Vertex
numVerticesSG :: Vertex
numEdgesSG :: Vertex
offsetSG :: Vector Vertex
adjacentSG :: Vector Vertex
edgeCtxSG :: Vector w
edgeCtxSG :: forall w. SparseGraph w -> Vector w
adjacentSG :: forall w. SparseGraph w -> Vector Vertex
offsetSG :: forall w. SparseGraph w -> Vector Vertex
numEdgesSG :: forall w. SparseGraph w -> Vertex
numVerticesSG :: forall w. SparseGraph w -> Vertex
..} = (forall s. ST s (MVector s Vertex)) -> Vector Vertex
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
U.create ((forall s. ST s (MVector s Vertex)) -> Vector Vertex)
-> (forall s. ST s (MVector s Vertex)) -> Vector Vertex
forall a b. (a -> b) -> a -> b
$ do
  dist <- Vertex -> Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
UM.replicate Vertex
numVerticesSG Vertex
forall a. Bounded a => a
maxBound
  que <- newBufferAsQueue (numEdgesSG + 1)
  UM.write dist source 0
  pushBack source que
  fix $ \ST s ()
loop -> do
    Buffer (PrimState (ST s)) Vertex -> ST s (Maybe Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Buffer (PrimState m) a -> m (Maybe a)
popFront Buffer s Vertex
Buffer (PrimState (ST s)) Vertex
que ST s (Maybe Vertex) -> (Maybe Vertex -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just Vertex
v -> do
        dv <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> m a
UM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
dist Vertex
v
        U.forM_ (gr `adj` v) $ \Vertex
nv -> do
          dnv <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> m a
UM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
dist Vertex
nv
          when (dnv == maxBound) $ do
            UM.unsafeWrite dist nv $ dv + 1
            pushBack nv que
        loop
      Maybe Vertex
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  return dist