module Data.Graph.Mutable
  ( -- * Graph Operations
    -- $mutgraph
    insertVertex
  , insertEdge
  , insertEdgeWith
    -- * Vertices Operations
    -- $mutvertices
  , verticesReplicate
  , verticesUReplicate
  , verticesWrite
  , verticesUWrite
  , verticesRead
  , verticesURead
  ) where

import Data.Graph.Types.Internal
import Control.Monad.Primitive
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed.Mutable as MU
import Data.Vector.Unboxed (Unbox)
import Data.Primitive.MutVar
import Data.Hashable (Hashable)
import qualified Data.HashMap.Mutable.Basic as HashTable

-- | $mutgraph
-- Operations that mutate a 'MGraph'. Vertices and edges can both be added,
-- and edges can be deleted, but vertices cannot be deleted. Providing such
-- an operation would undermine the safety that this library provides.

-- | This does two things:
--
--   * Check to see if a vertex with the provided value already exists
--   * Create a new vertex if it does not exist
--
--   In either case, the vertex id is returned, regardless or whether it was
--   preexisting or newly created.
insertVertex :: (PrimMonad m, Hashable v, Eq v) => MGraph (PrimState m) g e v -> v -> m (Vertex g)
insertVertex (MGraph vertexIndex currentIdVar _) v = do
  m <- HashTable.lookup vertexIndex v
  case m of
    Nothing -> do
      currentId <- readMutVar currentIdVar
      writeMutVar currentIdVar (currentId + 1)
      HashTable.insert vertexIndex v currentId
      return (Vertex currentId)
    Just i -> return (Vertex i)

-- | This replaces the edge if it already exists. If you pass the same vertex
--   as the source and the destination, this function has no effect.
insertEdge :: PrimMonad m => MGraph (PrimState m) g e v -> Vertex g -> Vertex g -> e -> m ()
insertEdge (MGraph _ _ edges) (Vertex a) (Vertex b) e = do
  HashTable.insert edges (IntPair a b) e

-- | Insert edge with a function, combining the existing edge value and the old one.
insertEdgeWith :: PrimMonad m => MGraph (PrimState m) g e v -> (e -> e -> e) -> Vertex g -> Vertex g -> e -> m ()
insertEdgeWith (MGraph _ _ edges) combine (Vertex a) (Vertex b) e = do
  m <- HashTable.lookup edges (IntPair a b)
  case m of
    Nothing -> HashTable.insert edges (IntPair a b) e
    Just eOld -> HashTable.insert edges (IntPair a b) (combine eOld e)

-- | $mutvertices
-- Operations that mutate a 'MVertices' or a 'MUVertices'. These functions have nothing
-- to do with 'MGraph' and are not usually needed by end users of this library. They
-- are useful for users writing algorithms that need to mark vertices in a graph as
-- it is traversed.
--
-- All of these operations are
-- wrappers around operations from @Data.Vector.Mutable@ and @Data.Vector.Unbox.Mutable@.
-- As long as you do not import @Data.Graph.Types.Internal@, this library guarentees that
-- these operations will not fail at runtime.

verticesReplicate :: PrimMonad m => Size g -> v -> m (MVertices (PrimState m) g v)
verticesReplicate (Size i) v = fmap MVertices (MV.replicate i v)

verticesUReplicate :: (PrimMonad m, Unbox v) => Size g -> v -> m (MUVertices (PrimState m) g v)
verticesUReplicate (Size i) v = fmap MUVertices (MU.replicate i v)

verticesUWrite :: (PrimMonad m, Unbox v) => MUVertices (PrimState m) g v -> Vertex g -> v -> m ()
verticesUWrite (MUVertices mvec) (Vertex ix) v = MU.unsafeWrite mvec ix v

verticesWrite :: PrimMonad m => MVertices (PrimState m) g v -> Vertex g -> v -> m ()
verticesWrite (MVertices mvec) (Vertex ix) v = MV.unsafeWrite mvec ix v

verticesURead :: (PrimMonad m, Unbox v) => MUVertices (PrimState m) g v -> Vertex g -> m v
verticesURead (MUVertices mvec) (Vertex ix) = MU.unsafeRead mvec ix

verticesRead :: PrimMonad m => MVertices (PrimState m) g v -> Vertex g -> m v
verticesRead (MVertices mvec) (Vertex ix) = MV.unsafeRead mvec ix