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

module Data.Graph.Sparse.BFS01 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

bfs01SG :: Vertex -> SparseGraph Int -> U.Vector Int
bfs01SG :: Int -> SparseGraph Int -> Vector Int
bfs01SG Int
source gr :: SparseGraph Int
gr@SparseGraph{Int
Vector Int
numVerticesSG :: Int
numEdgesSG :: Int
offsetSG :: Vector Int
adjacentSG :: Vector Int
edgeCtxSG :: Vector Int
edgeCtxSG :: forall w. SparseGraph w -> Vector w
adjacentSG :: forall w. SparseGraph w -> Vector Int
offsetSG :: forall w. SparseGraph w -> Vector Int
numEdgesSG :: forall w. SparseGraph w -> Int
numVerticesSG :: forall w. SparseGraph w -> Int
..} = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
U.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
  dist <- 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
numVerticesSG Int
forall a. Bounded a => a
maxBound
  deque <- newBufferAsDeque (numEdgesSG + 1)
  UM.write dist source 0
  pushFront (0, source) deque
  fix $ \ST s ()
loop ->
    Buffer (PrimState (ST s)) (Int, Int) -> ST s (Maybe (Int, Int))
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Buffer (PrimState m) a -> m (Maybe a)
popFront Buffer s (Int, Int)
Buffer (PrimState (ST s)) (Int, Int)
deque ST s (Maybe (Int, Int)) -> (Maybe (Int, Int) -> 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 (Int
dv, Int
v) -> do
        dv' <- 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
dist Int
v
        when (dv == dv') $ do
          U.forM_ (gr `adjW` v) $ \(Int
nv, Int
w) -> do
            dnv <- 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
dist Int
nv
            when (dv + w < dnv) $ do
              UM.unsafeWrite dist nv (dv + w)
              if w == 0
                then pushFront (dv, nv) deque
                else pushBack (dv + 1, nv) deque
        loop
      Maybe (Int, Int)
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  return dist