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