{-# LANGUAGE TypeFamilies #-}
-- | This graph is an efficient representation of bidirectional graphs with
-- parallel edges.
--
-- This is in contrast to 'Data.Graph.Haggle.SimpleBiDigraph', which
-- can only handle simple graphs (i.e., without parallel edges).
--
-- The representation is slightly less efficient as a result.
module Data.Graph.Haggle.BiDigraph (
  MBiDigraph,
  BiDigraph,
  newMBiDigraph,
  newSizedMBiDigraph
  ) where

import Control.Monad ( when )
import qualified Control.Monad.Primitive as P
import qualified Control.Monad.Ref as R
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IM
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector as V

import Data.Graph.Haggle.Classes
import Data.Graph.Haggle.Internal.Basic

-- | A mutable bidirectional graph
data MBiDigraph m =
  MBiDigraph { MBiDigraph m -> Ref m Int
mgraphVertexCount :: R.Ref m Int
             , MBiDigraph m -> Ref m Int
mgraphEdgeCount :: R.Ref m Int
             , MBiDigraph m -> Ref m Int
mgraphEdgeIdSrc :: R.Ref m Int
             , MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds :: R.Ref m (MV.MVector (P.PrimState m) (IntMap [Edge]))
             , MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs :: R.Ref m (MV.MVector (P.PrimState m) (IntMap [Edge]))
             }

-- | An immutable bidirectional graph
data BiDigraph =
  BiDigraph { BiDigraph -> Int
vertexCount :: {-# UNPACK #-} !Int
            , BiDigraph -> Int
edgeCount :: {-# UNPACK #-} !Int
            , BiDigraph -> Int
edgeIdSrc :: {-# UNPACK #-} !Int
            , BiDigraph -> Vector (IntMap [Edge])
graphPreds :: V.Vector (IntMap [Edge])
            , BiDigraph -> Vector (IntMap [Edge])
graphSuccs :: V.Vector (IntMap [Edge])
            }


defaultSize :: Int
defaultSize :: Int
defaultSize = Int
128

-- | Allocate a new mutable bidirectional graph with a default size
newMBiDigraph :: (P.PrimMonad m, R.MonadRef m) => m (MBiDigraph m)
newMBiDigraph :: m (MBiDigraph m)
newMBiDigraph = Int -> Int -> m (MBiDigraph m)
forall (m :: * -> *).
(PrimMonad m, MonadRef m) =>
Int -> Int -> m (MBiDigraph m)
newSizedMBiDigraph Int
defaultSize Int
0

-- | Allocate a new mutable bidirectional graph with space reserved
-- for nodes and edges.  This can be more efficient and avoid resizing.
newSizedMBiDigraph :: (P.PrimMonad m, R.MonadRef m)
                   => Int -- ^ Reserved space for nodes
                   -> Int -- ^ Reserved space for edges
                   -> m (MBiDigraph m)
newSizedMBiDigraph :: Int -> Int -> m (MBiDigraph m)
newSizedMBiDigraph Int
szNodes Int
_ = do
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
szNodes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"newSizedMBiDigraph: Negative size"
  Ref m Int
nn <- Int -> m (Ref m Int)
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef Int
0
  Ref m Int
en <- Int -> m (Ref m Int)
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef Int
0
  Ref m Int
esrc <- Int -> m (Ref m Int)
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef Int
0
  MVector (PrimState m) (IntMap [Edge])
pvec <- Int -> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.new Int
szNodes
  MVector (PrimState m) (IntMap [Edge])
svec <- Int -> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.new Int
szNodes
  Ref m (MVector (PrimState m) (IntMap [Edge]))
pref <- MVector (PrimState m) (IntMap [Edge])
-> m (Ref m (MVector (PrimState m) (IntMap [Edge])))
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef MVector (PrimState m) (IntMap [Edge])
pvec
  Ref m (MVector (PrimState m) (IntMap [Edge]))
sref <- MVector (PrimState m) (IntMap [Edge])
-> m (Ref m (MVector (PrimState m) (IntMap [Edge])))
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef MVector (PrimState m) (IntMap [Edge])
svec
  MBiDigraph m -> m (MBiDigraph m)
forall (m :: * -> *) a. Monad m => a -> m a
return (MBiDigraph m -> m (MBiDigraph m))
-> MBiDigraph m -> m (MBiDigraph m)
forall a b. (a -> b) -> a -> b
$! MBiDigraph :: forall (m :: * -> *).
Ref m Int
-> Ref m Int
-> Ref m Int
-> Ref m (MVector (PrimState m) (IntMap [Edge]))
-> Ref m (MVector (PrimState m) (IntMap [Edge]))
-> MBiDigraph m
MBiDigraph { mgraphVertexCount :: Ref m Int
mgraphVertexCount = Ref m Int
nn
                       , mgraphEdgeCount :: Ref m Int
mgraphEdgeCount = Ref m Int
en
                       , mgraphEdgeIdSrc :: Ref m Int
mgraphEdgeIdSrc = Ref m Int
esrc
                       , mgraphPreds :: Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds = Ref m (MVector (PrimState m) (IntMap [Edge]))
pref
                       , mgraphSuccs :: Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs = Ref m (MVector (PrimState m) (IntMap [Edge]))
sref
                       }

instance MGraph MBiDigraph where
  type ImmutableGraph MBiDigraph = BiDigraph
  getVertices :: MBiDigraph m -> m [Vertex]
getVertices MBiDigraph m
g = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    [Vertex] -> m [Vertex]
forall (m :: * -> *) a. Monad m => a -> m a
return [ Int -> Vertex
V Int
v | Int
v <- [Int
0.. Int
nVerts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ]

  getOutEdges :: MBiDigraph m -> Vertex -> m [Edge]
getOutEdges MBiDigraph m
g (V Int
src) = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    case Int
src Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nVerts of
      Bool
True -> [Edge] -> m [Edge]
forall (m :: * -> *) a. Monad m => a -> m a
return []
      Bool
False -> do
        MVector (PrimState m) (IntMap [Edge])
svec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g)
        IntMap [Edge]
succs <- MVector (PrimState m) (IntMap [Edge]) -> Int -> m (IntMap [Edge])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (IntMap [Edge])
svec Int
src
        [Edge] -> m [Edge]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Edge] -> m [Edge]) -> [Edge] -> m [Edge]
forall a b. (a -> b) -> a -> b
$ [[Edge]] -> [Edge]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (IntMap [Edge] -> [[Edge]]
forall a. IntMap a -> [a]
IM.elems IntMap [Edge]
succs)
  countVertices :: MBiDigraph m -> m Int
countVertices = Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (Ref m Int -> m Int)
-> (MBiDigraph m -> Ref m Int) -> MBiDigraph m -> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount
  countEdges :: MBiDigraph m -> m Int
countEdges = Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (Ref m Int -> m Int)
-> (MBiDigraph m -> Ref m Int) -> MBiDigraph m -> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphEdgeCount

  getSuccessors :: MBiDigraph m -> Vertex -> m [Vertex]
getSuccessors MBiDigraph m
g (V Int
src) = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    case Int
src Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nVerts of
      Bool
True -> [Vertex] -> m [Vertex]
forall (m :: * -> *) a. Monad m => a -> m a
return []
      Bool
False -> do
        MVector (PrimState m) (IntMap [Edge])
svec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g)
        IntMap [Edge]
succs <- MVector (PrimState m) (IntMap [Edge]) -> Int -> m (IntMap [Edge])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (IntMap [Edge])
svec Int
src
        [Vertex] -> m [Vertex]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Vertex] -> m [Vertex]) -> [Vertex] -> m [Vertex]
forall a b. (a -> b) -> a -> b
$ (Int -> Vertex) -> [Int] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Vertex
V ([Int] -> [Vertex]) -> [Int] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ IntMap [Edge] -> [Int]
forall a. IntMap a -> [Int]
IM.keys IntMap [Edge]
succs

  checkEdgeExists :: MBiDigraph m -> Vertex -> Vertex -> m Bool
checkEdgeExists MBiDigraph m
g (V Int
src) (V Int
dst) = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    case Int
src Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nVerts Bool -> Bool -> Bool
|| Int
dst Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nVerts of
      Bool
True -> Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
      Bool
False -> do
        MVector (PrimState m) (IntMap [Edge])
svec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g)
        IntMap [Edge]
succs <- MVector (PrimState m) (IntMap [Edge]) -> Int -> m (IntMap [Edge])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (IntMap [Edge])
svec Int
src
        Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Int -> IntMap [Edge] -> Bool
forall a. Int -> IntMap a -> Bool
IM.member Int
dst IntMap [Edge]
succs

  freeze :: MBiDigraph m -> m (ImmutableGraph MBiDigraph)
freeze MBiDigraph m
g = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    Int
nEdges <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphEdgeCount MBiDigraph m
g)
    Int
esrc <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphEdgeIdSrc MBiDigraph m
g)
    MVector (PrimState m) (IntMap [Edge])
pvec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds MBiDigraph m
g)
    MVector (PrimState m) (IntMap [Edge])
svec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g)
    Vector (IntMap [Edge])
pvec' <- MVector (PrimState m) (IntMap [Edge]) -> m (Vector (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze (Int
-> MVector (PrimState m) (IntMap [Edge])
-> MVector (PrimState m) (IntMap [Edge])
forall s a. Int -> MVector s a -> MVector s a
MV.take Int
nVerts MVector (PrimState m) (IntMap [Edge])
pvec)
    Vector (IntMap [Edge])
svec' <- MVector (PrimState m) (IntMap [Edge]) -> m (Vector (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze (Int
-> MVector (PrimState m) (IntMap [Edge])
-> MVector (PrimState m) (IntMap [Edge])
forall s a. Int -> MVector s a -> MVector s a
MV.take Int
nVerts MVector (PrimState m) (IntMap [Edge])
svec)
    BiDigraph -> m BiDigraph
forall (m :: * -> *) a. Monad m => a -> m a
return (BiDigraph -> m BiDigraph) -> BiDigraph -> m BiDigraph
forall a b. (a -> b) -> a -> b
$! BiDigraph :: Int
-> Int
-> Int
-> Vector (IntMap [Edge])
-> Vector (IntMap [Edge])
-> BiDigraph
BiDigraph { vertexCount :: Int
vertexCount = Int
nVerts
                        , edgeCount :: Int
edgeCount = Int
nEdges
                        , edgeIdSrc :: Int
edgeIdSrc = Int
esrc
                        , graphPreds :: Vector (IntMap [Edge])
graphPreds = Vector (IntMap [Edge])
pvec'
                        , graphSuccs :: Vector (IntMap [Edge])
graphSuccs = Vector (IntMap [Edge])
svec'
                        }

instance MAddVertex MBiDigraph where
  addVertex :: MBiDigraph m -> m Vertex
addVertex MBiDigraph m
g = do
    MBiDigraph m -> m ()
forall (m :: * -> *).
(PrimMonad m, MonadRef m) =>
MBiDigraph m -> m ()
ensureNodeSpace MBiDigraph m
g
    Int
vid <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef Ref m Int
r
    Ref m Int -> (Int -> Int) -> m ()
forall (m :: * -> *) a. MonadRef m => Ref m a -> (a -> a) -> m ()
R.modifyRef' Ref m Int
r (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
    MVector (PrimState m) (IntMap [Edge])
pvec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds MBiDigraph m
g)
    MVector (PrimState m) (IntMap [Edge])
svec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g)
    MVector (PrimState m) (IntMap [Edge])
-> Int -> IntMap [Edge] -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (IntMap [Edge])
pvec Int
vid IntMap [Edge]
forall a. IntMap a
IM.empty
    MVector (PrimState m) (IntMap [Edge])
-> Int -> IntMap [Edge] -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (IntMap [Edge])
svec Int
vid IntMap [Edge]
forall a. IntMap a
IM.empty
    Vertex -> m Vertex
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Vertex
V Int
vid)
    where
      r :: Ref m Int
r = MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g


instance MAddEdge MBiDigraph where
  addEdge :: MBiDigraph m -> Vertex -> Vertex -> m (Maybe Edge)
addEdge MBiDigraph m
g (V Int
src) (V Int
dst) = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    case Int
src Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nVerts Bool -> Bool -> Bool
|| Int
dst Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nVerts of
      Bool
True -> Maybe Edge -> m (Maybe Edge)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Edge
forall a. Maybe a
Nothing
      Bool
False -> do
        Int
eid <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphEdgeIdSrc MBiDigraph m
g)
        Ref m Int -> (Int -> Int) -> m ()
forall (m :: * -> *) a. MonadRef m => Ref m a -> (a -> a) -> m ()
R.modifyRef' (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphEdgeIdSrc MBiDigraph m
g) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
        Ref m Int -> (Int -> Int) -> m ()
forall (m :: * -> *) a. MonadRef m => Ref m a -> (a -> a) -> m ()
R.modifyRef' (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphEdgeCount MBiDigraph m
g) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
        let e :: Edge
e = Int -> Int -> Int -> Edge
E Int
eid Int
src Int
dst
        MVector (PrimState m) (IntMap [Edge])
pvec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds MBiDigraph m
g)
        IntMap [Edge]
preds <- MVector (PrimState m) (IntMap [Edge]) -> Int -> m (IntMap [Edge])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (IntMap [Edge])
pvec Int
dst
        MVector (PrimState m) (IntMap [Edge])
-> Int -> IntMap [Edge] -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector (PrimState m) (IntMap [Edge])
pvec Int
dst (([Edge] -> [Edge] -> [Edge])
-> Int -> [Edge] -> IntMap [Edge] -> IntMap [Edge]
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IM.insertWith [Edge] -> [Edge] -> [Edge]
forall a. [a] -> [a] -> [a]
(++) Int
src [Edge
e] IntMap [Edge]
preds)

        MVector (PrimState m) (IntMap [Edge])
svec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g)
        IntMap [Edge]
succs <- MVector (PrimState m) (IntMap [Edge]) -> Int -> m (IntMap [Edge])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (IntMap [Edge])
svec Int
src
        MVector (PrimState m) (IntMap [Edge])
-> Int -> IntMap [Edge] -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector (PrimState m) (IntMap [Edge])
svec Int
src (([Edge] -> [Edge] -> [Edge])
-> Int -> [Edge] -> IntMap [Edge] -> IntMap [Edge]
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IM.insertWith [Edge] -> [Edge] -> [Edge]
forall a. [a] -> [a] -> [a]
(++) Int
dst [Edge
e] IntMap [Edge]
succs)

        Maybe Edge -> m (Maybe Edge)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Edge -> m (Maybe Edge)) -> Maybe Edge -> m (Maybe Edge)
forall a b. (a -> b) -> a -> b
$ Edge -> Maybe Edge
forall a. a -> Maybe a
Just Edge
e

instance MBidirectional MBiDigraph where
  getPredecessors :: MBiDigraph m -> Vertex -> m [Vertex]
getPredecessors MBiDigraph m
g (V Int
vid) = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    case Int
vid Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nVerts of
      Bool
False -> [Vertex] -> m [Vertex]
forall (m :: * -> *) a. Monad m => a -> m a
return []
      Bool
True -> do
        MVector (PrimState m) (IntMap [Edge])
pvec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds MBiDigraph m
g)
        IntMap [Edge]
preds <- MVector (PrimState m) (IntMap [Edge]) -> Int -> m (IntMap [Edge])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (IntMap [Edge])
pvec Int
vid
        [Vertex] -> m [Vertex]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Vertex] -> m [Vertex]) -> [Vertex] -> m [Vertex]
forall a b. (a -> b) -> a -> b
$ (Int -> Vertex) -> [Int] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Vertex
V ([Int] -> [Vertex]) -> [Int] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ IntMap [Edge] -> [Int]
forall a. IntMap a -> [Int]
IM.keys IntMap [Edge]
preds

  getInEdges :: MBiDigraph m -> Vertex -> m [Edge]
getInEdges MBiDigraph m
g (V Int
vid) = do
    Int
nVerts <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
    case Int
vid Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nVerts of
      Bool
False -> [Edge] -> m [Edge]
forall (m :: * -> *) a. Monad m => a -> m a
return []
      Bool
True -> do
        MVector (PrimState m) (IntMap [Edge])
pvec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds MBiDigraph m
g)
        IntMap [Edge]
preds <- MVector (PrimState m) (IntMap [Edge]) -> Int -> m (IntMap [Edge])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (IntMap [Edge])
pvec Int
vid
        [Edge] -> m [Edge]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Edge] -> m [Edge]) -> [Edge] -> m [Edge]
forall a b. (a -> b) -> a -> b
$ [[Edge]] -> [Edge]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (IntMap [Edge] -> [[Edge]]
forall a. IntMap a -> [a]
IM.elems IntMap [Edge]
preds)

instance Thawable BiDigraph where
  type MutableGraph BiDigraph = MBiDigraph
  thaw :: BiDigraph -> m (MutableGraph BiDigraph m)
thaw BiDigraph
g = do
    Ref m Int
vc <- Int -> m (Ref m Int)
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef (BiDigraph -> Int
vertexCount BiDigraph
g)
    Ref m Int
ec <- Int -> m (Ref m Int)
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef (BiDigraph -> Int
edgeCount BiDigraph
g)
    Ref m Int
eidsrc <- Int -> m (Ref m Int)
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef (BiDigraph -> Int
edgeIdSrc BiDigraph
g)
    MVector (PrimState m) (IntMap [Edge])
pvec <- Vector (IntMap [Edge]) -> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.thaw (BiDigraph -> Vector (IntMap [Edge])
graphPreds BiDigraph
g)
    MVector (PrimState m) (IntMap [Edge])
svec <- Vector (IntMap [Edge]) -> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.thaw (BiDigraph -> Vector (IntMap [Edge])
graphSuccs BiDigraph
g)
    Ref m (MVector (PrimState m) (IntMap [Edge]))
pref <- MVector (PrimState m) (IntMap [Edge])
-> m (Ref m (MVector (PrimState m) (IntMap [Edge])))
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef MVector (PrimState m) (IntMap [Edge])
pvec
    Ref m (MVector (PrimState m) (IntMap [Edge]))
sref <- MVector (PrimState m) (IntMap [Edge])
-> m (Ref m (MVector (PrimState m) (IntMap [Edge])))
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef MVector (PrimState m) (IntMap [Edge])
svec
    MBiDigraph m -> m (MBiDigraph m)
forall (m :: * -> *) a. Monad m => a -> m a
return MBiDigraph :: forall (m :: * -> *).
Ref m Int
-> Ref m Int
-> Ref m Int
-> Ref m (MVector (PrimState m) (IntMap [Edge]))
-> Ref m (MVector (PrimState m) (IntMap [Edge]))
-> MBiDigraph m
MBiDigraph { mgraphVertexCount :: Ref m Int
mgraphVertexCount = Ref m Int
vc
                      , mgraphEdgeCount :: Ref m Int
mgraphEdgeCount = Ref m Int
ec
                      , mgraphEdgeIdSrc :: Ref m Int
mgraphEdgeIdSrc = Ref m Int
eidsrc
                      , mgraphPreds :: Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds = Ref m (MVector (PrimState m) (IntMap [Edge]))
pref
                      , mgraphSuccs :: Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs = Ref m (MVector (PrimState m) (IntMap [Edge]))
sref
                      }

instance Graph BiDigraph where
  vertices :: BiDigraph -> [Vertex]
vertices BiDigraph
g = (Int -> Vertex) -> [Int] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Vertex
V [Int
0 .. BiDigraph -> Int
vertexCount BiDigraph
g Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
  edges :: BiDigraph -> [Edge]
edges BiDigraph
g = (Vertex -> [Edge]) -> [Vertex] -> [Edge]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (BiDigraph -> Vertex -> [Edge]
forall g. Graph g => g -> Vertex -> [Edge]
outEdges BiDigraph
g) (BiDigraph -> [Vertex]
forall g. Graph g => g -> [Vertex]
vertices BiDigraph
g)
  successors :: BiDigraph -> Vertex -> [Vertex]
successors BiDigraph
g (V Int
v)
    | BiDigraph -> Int -> Bool
outOfRange BiDigraph
g Int
v = []
    | Bool
otherwise = (Int -> Vertex) -> [Int] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Vertex
V ([Int] -> [Vertex]) -> [Int] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ IntMap [Edge] -> [Int]
forall a. IntMap a -> [Int]
IM.keys (IntMap [Edge] -> [Int]) -> IntMap [Edge] -> [Int]
forall a b. (a -> b) -> a -> b
$ Vector (IntMap [Edge]) -> Int -> IntMap [Edge]
forall a. Vector a -> Int -> a
V.unsafeIndex (BiDigraph -> Vector (IntMap [Edge])
graphSuccs BiDigraph
g) Int
v
  outEdges :: BiDigraph -> Vertex -> [Edge]
outEdges BiDigraph
g (V Int
v)
    | BiDigraph -> Int -> Bool
outOfRange BiDigraph
g Int
v = []
    | Bool
otherwise =
      let succs :: IntMap [Edge]
succs = Vector (IntMap [Edge]) -> Int -> IntMap [Edge]
forall a. Vector a -> Int -> a
V.unsafeIndex (BiDigraph -> Vector (IntMap [Edge])
graphSuccs BiDigraph
g) Int
v
      in [[Edge]] -> [Edge]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (IntMap [Edge] -> [[Edge]]
forall a. IntMap a -> [a]
IM.elems IntMap [Edge]
succs)
  edgesBetween :: BiDigraph -> Vertex -> Vertex -> [Edge]
edgesBetween BiDigraph
g (V Int
src) (V Int
dst)
    | BiDigraph -> Int -> Bool
outOfRange BiDigraph
g Int
src Bool -> Bool -> Bool
|| BiDigraph -> Int -> Bool
outOfRange BiDigraph
g Int
dst = []
    | Bool
otherwise = [Edge] -> Int -> IntMap [Edge] -> [Edge]
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault [] Int
dst (Vector (IntMap [Edge]) -> Int -> IntMap [Edge]
forall a. Vector a -> Int -> a
V.unsafeIndex (BiDigraph -> Vector (IntMap [Edge])
graphSuccs BiDigraph
g) Int
src)
  maxVertexId :: BiDigraph -> Int
maxVertexId BiDigraph
g = Vector (IntMap [Edge]) -> Int
forall a. Vector a -> Int
V.length (BiDigraph -> Vector (IntMap [Edge])
graphSuccs BiDigraph
g) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  isEmpty :: BiDigraph -> Bool
isEmpty = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0) (Int -> Bool) -> (BiDigraph -> Int) -> BiDigraph -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiDigraph -> Int
vertexCount

instance Bidirectional BiDigraph  where
  predecessors :: BiDigraph -> Vertex -> [Vertex]
predecessors BiDigraph
g (V Int
v)
    | BiDigraph -> Int -> Bool
outOfRange BiDigraph
g Int
v = []
    | Bool
otherwise = (Int -> Vertex) -> [Int] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Vertex
V ([Int] -> [Vertex]) -> [Int] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ IntMap [Edge] -> [Int]
forall a. IntMap a -> [Int]
IM.keys (IntMap [Edge] -> [Int]) -> IntMap [Edge] -> [Int]
forall a b. (a -> b) -> a -> b
$ Vector (IntMap [Edge]) -> Int -> IntMap [Edge]
forall a. Vector a -> Int -> a
V.unsafeIndex (BiDigraph -> Vector (IntMap [Edge])
graphPreds BiDigraph
g) Int
v
  inEdges :: BiDigraph -> Vertex -> [Edge]
inEdges BiDigraph
g (V Int
v)
    | BiDigraph -> Int -> Bool
outOfRange BiDigraph
g Int
v = []
    | Bool
otherwise =
      let preds :: IntMap [Edge]
preds = Vector (IntMap [Edge]) -> Int -> IntMap [Edge]
forall a. Vector a -> Int -> a
V.unsafeIndex (BiDigraph -> Vector (IntMap [Edge])
graphPreds BiDigraph
g) Int
v
      in [[Edge]] -> [Edge]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (IntMap [Edge] -> [[Edge]]
forall a. IntMap a -> [a]
IM.elems IntMap [Edge]
preds)

-- Helpers

outOfRange :: BiDigraph -> Int -> Bool
outOfRange :: BiDigraph -> Int -> Bool
outOfRange BiDigraph
g = (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= BiDigraph -> Int
vertexCount BiDigraph
g)

-- | Given a graph, ensure that there is space in the vertex vector
-- for a new vertex.  If there is not, double the capacity.
ensureNodeSpace :: (P.PrimMonad m, R.MonadRef m) => MBiDigraph m -> m ()
ensureNodeSpace :: MBiDigraph m -> m ()
ensureNodeSpace MBiDigraph m
g = do
  MVector (PrimState m) (IntMap [Edge])
pvec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds MBiDigraph m
g)
  MVector (PrimState m) (IntMap [Edge])
svec <- Ref m (MVector (PrimState m) (IntMap [Edge]))
-> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g)
  let cap :: Int
cap = MVector (PrimState m) (IntMap [Edge]) -> Int
forall s a. MVector s a -> Int
MV.length MVector (PrimState m) (IntMap [Edge])
pvec
  Int
cnt <- Ref m Int -> m Int
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef (MBiDigraph m -> Ref m Int
forall (m :: * -> *). MBiDigraph m -> Ref m Int
mgraphVertexCount MBiDigraph m
g)
  case Int
cnt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
cap of
    Bool
True -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Bool
False -> do
      MVector (PrimState m) (IntMap [Edge])
pvec' <- MVector (PrimState m) (IntMap [Edge])
-> Int -> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
MV.grow MVector (PrimState m) (IntMap [Edge])
pvec Int
cap
      MVector (PrimState m) (IntMap [Edge])
svec' <- MVector (PrimState m) (IntMap [Edge])
-> Int -> m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
MV.grow MVector (PrimState m) (IntMap [Edge])
svec Int
cap
      Ref m (MVector (PrimState m) (IntMap [Edge]))
-> MVector (PrimState m) (IntMap [Edge]) -> m ()
forall (m :: * -> *) a. MonadRef m => Ref m a -> a -> m ()
R.writeRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphPreds MBiDigraph m
g) MVector (PrimState m) (IntMap [Edge])
pvec'
      Ref m (MVector (PrimState m) (IntMap [Edge]))
-> MVector (PrimState m) (IntMap [Edge]) -> m ()
forall (m :: * -> *) a. MonadRef m => Ref m a -> a -> m ()
R.writeRef (MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
forall (m :: * -> *).
MBiDigraph m -> Ref m (MVector (PrimState m) (IntMap [Edge]))
mgraphSuccs MBiDigraph m
g) MVector (PrimState m) (IntMap [Edge])
svec'