{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE RankNTypes #-}
module Data.Graph.Immutable
(
lookupVertex
, lookupEdge
, atVertex
, mapVertices
, mapEdges
, traverseVertices_
, traverseEdges_
, traverseNeighbors_
, vertices
, setVertices
, size
, freeze
, create
, with
, mapSome
, dijkstra
, dijkstraDistance
, dijkstraFoldM
, sizeInt
, vertexInt
, verticesRead
, verticesLength
, verticesTraverse
, verticesTraverse_
, verticesToVertexList
, verticesToVector
, verticesThaw
, verticesFreeze
) where
import Data.Graph.Types.Internal
import Control.Monad.Primitive
import Data.Foldable
import Data.Vector (Vector)
import Data.Vector.Mutable (MVector)
import Data.Functor.Identity (Identity(..))
import Control.Monad
import Data.Word
import Control.Monad.ST (runST)
import Data.Primitive.MutVar
import Data.Coerce (coerce)
import Data.Semigroup (Semigroup)
import qualified Data.Graph.Mutable as Mutable
import qualified Data.ArrayList.Generic as ArrayList
import qualified Data.HashMap.Mutable.Basic as HashTable
import qualified Data.Heap.Mutable.ModelD as Heap
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import qualified Data.Semigroup as SG
lookupVertex :: Eq v => v -> Graph g e v -> Maybe (Vertex g)
lookupVertex val (Graph g) = fmap Vertex (V.elemIndex val (graphVertices g))
lookupEdge :: Vertex g -> Vertex g -> Graph g e v -> Maybe e
lookupEdge (Vertex x) (Vertex y) (Graph (SomeGraph _ neighbors edges)) =
case U.elemIndex y (V.unsafeIndex neighbors x) of
Nothing -> Nothing
Just ix -> Just (V.unsafeIndex (V.unsafeIndex edges x) ix)
atVertex :: Vertex g -> Graph g e v -> v
atVertex v g = verticesRead (vertices g) v
mapVertices :: (Vertex g -> a -> b) -> Graph g e a -> Graph g e b
mapVertices f (Graph sg) = Graph sg
{ graphVertices = V.imap (coerce f) (graphVertices sg) }
mapEdges :: (Vertex g -> Vertex g -> e -> d) -> Graph g e v -> Graph g d v
mapEdges f (Graph (SomeGraph v verts edges)) = Graph $ SomeGraph v verts $
V.imap
( \outerIx edgeVals ->
let vertIxs = V.unsafeIndex verts outerIx
in V.imap
( \sourceIx edgeVal ->
let destIx = U.unsafeIndex vertIxs sourceIx
in f (Vertex sourceIx) (Vertex destIx) edgeVal
) edgeVals
) edges
traverseVertices_ :: Applicative m => (Vertex g -> v -> m a) -> Graph g e v -> m ()
traverseVertices_ f g = verticesTraverse_ f (vertices g)
traverseEdges_ :: Applicative m
=> (Vertex g -> Vertex g -> v -> v -> e -> m a)
-> Graph g e v
-> m ()
traverseEdges_ f g =
let allVertices = vertices g
in verticesTraverse_
(\vertex value -> traverseNeighbors_
(\neighborVertex neighborValue e -> f vertex neighborVertex value neighborValue e)
vertex g
) allVertices
traverseNeighbors_ :: Applicative m
=> (Vertex g -> v -> e -> m a)
-> Vertex g
-> Graph g e v
-> m ()
traverseNeighbors_ f (Vertex x) (Graph g) =
let allVertices = graphVertices g
theVertices = V.unsafeIndex (graphOutboundNeighborVertices g) x
edges = V.unsafeIndex (graphOutboundNeighborEdges g) x
numNeighbors = U.length theVertices
go !i = if i < numNeighbors
then let vertexNum = U.unsafeIndex theVertices i
vertexVal = V.unsafeIndex allVertices vertexNum
edgeVal = V.unsafeIndex edges i
in f (Vertex vertexNum) vertexVal edgeVal *> go (i + 1)
else pure ()
in go 0
vertices :: Graph g e v -> Vertices g v
vertices (Graph (SomeGraph v _ _)) = Vertices v
setVertices :: Vertices g v -> Graph g e w -> Graph g e v
setVertices (Vertices x) (Graph (SomeGraph _ a b)) = Graph (SomeGraph x a b)
size :: Graph g e v -> Size g
size (Graph (SomeGraph v _ _)) = Size (V.length v)
sizeInt :: Size g -> Int
sizeInt (Size s) = s
vertexInt :: Vertex g -> Int
vertexInt (Vertex i) = i
verticesToVertexList :: Vertices g v -> [Vertex g]
verticesToVertexList (Vertices v) = map Vertex (take (V.length v) [0..])
verticesTraverse :: Applicative m => (Vertex g -> v -> m a) -> Vertices g v -> m (Vertices g a)
verticesTraverse f (Vertices v) = fmap (Vertices . V.fromList) $ traverse (\(i,b) -> f (Vertex i) b) (zip [0..] (V.toList v))
verticesTraverse_ :: Applicative m => (Vertex g -> v -> m a) -> Vertices g v -> m ()
verticesTraverse_ f (Vertices v) = traverse_ (\(i,b) -> f (Vertex i) b) (zip [0..] (V.toList v))
verticesToVector :: Vertices g v -> Vector v
verticesToVector (Vertices v) = v
verticesRead :: Vertices g v -> Vertex g -> v
verticesRead (Vertices v) (Vertex i) = V.unsafeIndex v i
verticesLength :: Vertices g v -> Int
verticesLength (Vertices v) = V.length v
verticesFreeze :: PrimMonad m => MVertices (PrimState m) g v -> m (Vertices g v)
verticesFreeze (MVertices mvec) = fmap Vertices (V.freeze mvec)
verticesThaw :: PrimMonad m => Vertices g v -> m (MVertices (PrimState m) g v)
verticesThaw (Vertices vec) = fmap MVertices (V.thaw vec)
freeze :: PrimMonad m => MGraph (PrimState m) g e v -> m (Graph g e v)
freeze (MGraph vertexIndex currentIdVar edges) = do
let initialArrayListSize = 16
numberOfVertices <- readMutVar currentIdVar
mvec <- MV.new numberOfVertices
mvecEdgeVals <- MV.replicateM numberOfVertices (ArrayList.new initialArrayListSize)
mvecEdgeNeighbors <- MV.replicateM numberOfVertices (ArrayList.new initialArrayListSize)
flip HashTable.mapM_ vertexIndex $ \vertexValue vertexId -> do
MV.unsafeWrite mvec vertexId vertexValue
flip HashTable.mapM_ edges $ \(IntPair fromVertexId toVertexId) edgeVal -> do
mvecEdgeVal <- MV.unsafeRead mvecEdgeVals fromVertexId
ArrayList.push mvecEdgeVal edgeVal
mvecEdgeNeighbor <- MV.unsafeRead mvecEdgeNeighbors fromVertexId
ArrayList.push mvecEdgeNeighbor toVertexId
vecEdgeVals1 <- V.unsafeFreeze mvecEdgeVals
vecEdgeVals2 <- V.mapM ArrayList.freeze vecEdgeVals1
vecEdgeNeighbors1 <- V.unsafeFreeze mvecEdgeNeighbors
vecEdgeNeighbors2 <- V.mapM ArrayList.freeze vecEdgeNeighbors1
vec <- V.unsafeFreeze mvec
return (Graph $ SomeGraph vec vecEdgeNeighbors2 vecEdgeVals2)
create :: PrimMonad m => (forall g. MGraph (PrimState m) g e v -> m ()) -> m (SomeGraph e v)
create f = do
mg <- MGraph
<$> HashTable.new
<*> newMutVar 0
<*> HashTable.new
f mg
Graph g <- freeze mg
return g
with :: SomeGraph e v -> (forall g. Graph g e v -> a) -> a
with sg f = f (Graph sg)
mapSome :: (forall g. Graph g e v -> Graph g e' v') -> SomeGraph e v -> SomeGraph e' v'
mapSome f g = case f (Graph g) of
Graph g' -> g'
dijkstraDistance :: (Num e, Ord e)
=> Vertex g
-> Vertex g
-> Graph g e v
-> Maybe e
dijkstraDistance start end g =
getMinDistance $ atVertex end
( dijkstra
(\_ _ mdist e -> addMinDistance mdist e)
(MinDistance (Just 0))
(Identity start) g
)
where addMinDistance (MinDistance m) e = MinDistance (fmap (+ e) m)
newtype MinDistance a = MinDistance { getMinDistance :: Maybe a }
instance Eq a => Eq (MinDistance a) where
MinDistance a == MinDistance b = a == b
instance Ord a => Ord (MinDistance a) where
compare (MinDistance a) (MinDistance b) = case a of
Nothing -> case b of
Nothing -> EQ
Just _ -> GT
Just aval -> case b of
Nothing -> LT
Just bval -> compare aval bval
instance Ord a => Semigroup (MinDistance a) where
(<>) = min
instance Ord a => Monoid (MinDistance a) where
mempty = MinDistance Nothing
mappend = (SG.<>)
dijkstra ::
(Ord s, Monoid s, Foldable t)
=> (v -> v -> s -> e -> s)
-> s
-> t (Vertex g)
-> Graph g e v
-> Graph g e s
dijkstra f s0 v0 g =
fst $ runST $ dijkstraGeneral f (\_ _ _ -> return ()) s0 () v0 g
dijkstraFoldM ::
(Ord s, Monoid s, Foldable t, PrimMonad m)
=> (v -> v -> s -> e -> s)
-> (v -> s -> x -> m x)
-> s
-> x
-> t (Vertex g)
-> Graph g e v
-> m x
dijkstraFoldM f mf s0 acc v0 g =
fmap snd $ dijkstraGeneral f mf s0 acc v0 g
dijkstraGeneral ::
(Ord s, Monoid s, Foldable t, PrimMonad m)
=> (v -> v -> s -> e -> s)
-> (v -> s -> x -> m x)
-> s
-> x
-> t (Vertex g)
-> Graph g e v
-> m (Graph g e s, x)
dijkstraGeneral f step s0 x0 v0 g = do
let theSize = size g
oldVertices = vertices g
newVertices <- Mutable.verticesReplicate theSize mempty
visited <- Mutable.verticesUReplicate theSize False
heap <- Heap.new (sizeInt theSize)
forM_ v0 $ \v -> do
Mutable.verticesWrite newVertices v s0
Heap.unsafePush s0 (getVertexInternal v) heap
let go x = do
m <- Heap.pop heap
case m of
Nothing -> return (BoolWith True x)
Just (s,unwrappedVertexIx) -> do
let vertex = Vertex unwrappedVertexIx
value = verticesRead oldVertices vertex
Mutable.verticesUWrite visited vertex True
Mutable.verticesWrite newVertices vertex s
traverseNeighbors_ (\neighborVertex neighborValue theEdge -> do
alreadyVisited <- Mutable.verticesURead visited neighborVertex
when (not alreadyVisited) $ Heap.unsafePush
(f value neighborValue s theEdge)
(getVertexInternal neighborVertex)
heap
) vertex g
xNext <- step value s x
return (BoolWith False xNext)
runMe x = do
BoolWith isDone xNext <- go x
if isDone then return xNext else runMe xNext
xFinal <- runMe x0
newVerticesFrozen <- verticesFreeze newVertices
return (setVertices newVerticesFrozen g, xFinal)
data BoolWith a = BoolWith Bool !a