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

module Data.Graph.Sparse.Dijkstra where

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

import Data.Graph.Sparse
import Data.Heap.Binary

dijkstraSG ::
  (U.Unbox w, Num w, Ord w, Bounded w) =>
  Vertex ->
  SparseGraph w ->
  U.Vector w
dijkstraSG :: forall w.
(Unbox w, Num w, Ord w, Bounded w) =>
Vertex -> SparseGraph w -> Vector w
dijkstraSG 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 w)) -> Vector w
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
U.create ((forall s. ST s (MVector s w)) -> Vector w)
-> (forall s. ST s (MVector s w)) -> Vector w
forall a b. (a -> b) -> a -> b
$ do
  dist <- Vertex -> w -> ST s (MVector (PrimState (ST s)) w)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
UM.replicate Vertex
numVerticesSG w
forall a. Bounded a => a
maxBound
  heap <- newMinBinaryHeap (numEdgesSG + 1)
  UM.write dist source 0
  insertBH (0, source) heap
  fix $ \ST s ()
loop -> do
    MinBinaryHeap (PrimState (ST s)) (w, Vertex)
-> ST s (Maybe (w, Vertex))
forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m (Maybe a)
deleteFindTopBH MinBinaryHeap s (w, Vertex)
MinBinaryHeap (PrimState (ST s)) (w, Vertex)
heap ST s (Maybe (w, Vertex))
-> (Maybe (w, 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 (w
d, Vertex
v) -> do
        dv <- MVector (PrimState (ST s)) w -> Vertex -> ST s w
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> m a
UM.unsafeRead MVector s w
MVector (PrimState (ST s)) w
dist Vertex
v
        when (dv == d) $ do
          U.forM_ (gr `adjW` v) $ \(Vertex
nv, w
w) -> do
            dnv <- MVector (PrimState (ST s)) w -> Vertex -> ST s w
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> m a
UM.unsafeRead MVector s w
MVector (PrimState (ST s)) w
dist Vertex
nv
            when (dv + w < dnv) $ do
              UM.unsafeWrite dist nv $ dv + w
              insertBH (dv + w, nv) heap
        loop
      Maybe (w, Vertex)
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  return dist