module Data.Graph.Immutable
(
lookupVertex
, lookupEdge
, mapVertices
, traverseVertices_
, traverseEdges_
, traverseNeighbors_
, vertices
, setVertices
, size
, freeze
, create
, with
, dijkstra
, dijkstraMonoidal
, dijkstraMonoidalCover
, sizeInt
, vertexInt
, verticesRead
, verticesLength
, 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 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
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)
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) }
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 = graphOutboundNeighborVertices g V.! x
edges = graphOutboundNeighborEdges g V.! x
numNeighbors = U.length theVertices
go !i = if i < numNeighbors
then let vertexNum = theVertices U.! i
vertexVal = allVertices V.! vertexNum
edgeVal = edges V.! 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 v -> 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)
dijkstra :: (Num e, Ord e)
=> Vertex g
-> Vertex g
-> Graph g e v
-> Maybe e
dijkstra start end g = getMinDistance
( dijkstraMonoidal
(\_ _ mdist e -> addMinDistance mdist e)
(MinDistance (Just 0))
start end 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 => Monoid (MinDistance a) where
mempty = MinDistance Nothing
mappend ma mb = min ma mb
dijkstraMonoidal :: (Ord s, Monoid s)
=> (v -> v -> s -> e -> s)
-> s
-> Vertex g
-> Vertex g
-> Graph g e v
-> s
dijkstraMonoidal f s start end g =
verticesRead (dijkstraMonoidalCover f s (Identity start) g) end
dijkstraMonoidalCover ::
(Ord s, Monoid s, Foldable t)
=> (v -> v -> s -> e -> s)
-> s
-> t (Vertex g)
-> Graph g e v
-> Vertices g s
dijkstraMonoidalCover f s0 v0 g = runST $ 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 = do
m <- Heap.pop heap
case m of
Nothing -> return True
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
return False
runMe = do
isDone <- go
if isDone then return () else runMe
runMe
newVerticesFrozen <- verticesFreeze newVertices
return newVerticesFrozen