{-# LANGUAGE PatternGuards, FlexibleContexts #-}
-- | This is a simple module to handle a common pattern: constructing graphs
-- where vertex labels map uniquely to vertices.
--
-- The primary functions in this module are 'vertexForLabel' and
-- 'vertexForLabelRef', which take a vertex label and return the 'Vertex' for
-- that label (allocating a new 'Vertex') if necessary.  The first of those
-- functions explicitly threads the mapping as inputs and outputs.  The second
-- manages a mutable ref side-by-side with the underlying graph.
--
-- After the graph is fully constructed, this mapping is often still useful.
module Data.Graph.Haggle.VertexMap (
  -- * Pure interface
  VertexMap,
  emptyVertexMap,
  vertexForLabel,
  lookupVertexForLabel,
  vertexMapFromGraph,
  -- * Ref interface
  VertexMapRef,
  newVertexMapRef,
  vertexForLabelRef,
  vertexMapFromRef ) where

import qualified Control.DeepSeq as DS
import Control.Monad ( liftM )
import qualified Control.Monad.Primitive as P
import qualified Control.Monad.Ref as R
import Data.Map ( Map )
import qualified Data.Map as M
import Data.Tuple ( swap )

import Data.Graph.Haggle.Classes

-- | A simple mapping from labels to their 'Vertex'
newtype VertexMap nl = VM (Map nl Vertex)

instance (DS.NFData nl) => DS.NFData (VertexMap nl) where
  rnf :: VertexMap nl -> ()
rnf (VM Map nl Vertex
m) = Map nl Vertex
m forall a b. NFData a => a -> b -> b
`DS.deepseq` ()

emptyVertexMap :: VertexMap nl
emptyVertexMap :: forall nl. VertexMap nl
emptyVertexMap = forall nl. Map nl Vertex -> VertexMap nl
VM forall k a. Map k a
M.empty

-- | > (v, m') <- vertexForLabel g m lbl
--
-- Looks up the 'Vertex' for @lbl@ in @g@.  If no 'Vertex' in @g@ has that
-- label, a new 'Vertex' is allocated and returned.  The updated vertex
-- mapping @m'@ is returned, too.
vertexForLabel :: (MLabeledVertex g, Ord (MVertexLabel g), P.PrimMonad m, R.MonadRef m)
               => g m
               -> VertexMap (MVertexLabel g)
               -> MVertexLabel g
               -> m (Vertex, VertexMap (MVertexLabel g))
vertexForLabel :: forall (g :: (* -> *) -> *) (m :: * -> *).
(MLabeledVertex g, Ord (MVertexLabel g), PrimMonad m,
 MonadRef m) =>
g m
-> VertexMap (MVertexLabel g)
-> MVertexLabel g
-> m (Vertex, VertexMap (MVertexLabel g))
vertexForLabel g m
g vm :: VertexMap (MVertexLabel g)
vm@(VM Map (MVertexLabel g) Vertex
m) MVertexLabel g
lbl
  | Just Vertex
v <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup MVertexLabel g
lbl Map (MVertexLabel g) Vertex
m = forall (m :: * -> *) a. Monad m => a -> m a
return (Vertex
v, VertexMap (MVertexLabel g)
vm)
  | Bool
otherwise = do
    Vertex
v <- forall (g :: (* -> *) -> *) (m :: * -> *).
(MLabeledVertex g, PrimMonad m, MonadRef m) =>
g m -> MVertexLabel g -> m Vertex
addLabeledVertex g m
g MVertexLabel g
lbl
    let m' :: Map (MVertexLabel g) Vertex
m' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert MVertexLabel g
lbl Vertex
v Map (MVertexLabel g) Vertex
m
    forall (m :: * -> *) a. Monad m => a -> m a
return (Vertex
v, forall nl. Map nl Vertex -> VertexMap nl
VM Map (MVertexLabel g) Vertex
m')

-- | A pure lookup to convert a 'Vertex' label into a 'Vertex'.  If the
-- label is not in the graph, returns 'Nothing'.
lookupVertexForLabel :: (Ord nl) => nl -> VertexMap nl -> Maybe Vertex
lookupVertexForLabel :: forall nl. Ord nl => nl -> VertexMap nl -> Maybe Vertex
lookupVertexForLabel nl
lbl (VM Map nl Vertex
m) = forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup nl
lbl Map nl Vertex
m

-- | Build a 'VertexMap' from a 'Graph' with 'Vertex' labels.
vertexMapFromGraph :: (HasVertexLabel g, Ord (VertexLabel g))
                   => g -> VertexMap (VertexLabel g)
vertexMapFromGraph :: forall g.
(HasVertexLabel g, Ord (VertexLabel g)) =>
g -> VertexMap (VertexLabel g)
vertexMapFromGraph = forall nl. Map nl Vertex -> VertexMap nl
VM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> (b, a)
swap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall g. HasVertexLabel g => g -> [(Vertex, VertexLabel g)]
labeledVertices

-- | A 'VertexMap' wrapped up in a mutable ref for possibly
-- easier access in 'vertexMapFromRef'.
newtype VertexMapRef nl m = VMR (R.Ref m (VertexMap nl))

-- | Extract the pure 'VertexMap' from the mutable ref.  This is useful
-- to retain the mapping after the graph is fully constructed.
vertexMapFromRef :: (P.PrimMonad m, R.MonadRef m) => VertexMapRef nl m -> m (VertexMap nl)
vertexMapFromRef :: forall (m :: * -> *) nl.
(PrimMonad m, MonadRef m) =>
VertexMapRef nl m -> m (VertexMap nl)
vertexMapFromRef (VMR Ref m (VertexMap nl)
ref) = forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef Ref m (VertexMap nl)
ref

-- | Allocate a new 'VertexMap' buried in a mutable ref.
newVertexMapRef :: (P.PrimMonad m, R.MonadRef m) => m (VertexMapRef nl m)
newVertexMapRef :: forall (m :: * -> *) nl.
(PrimMonad m, MonadRef m) =>
m (VertexMapRef nl m)
newVertexMapRef = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall nl (m :: * -> *). Ref m (VertexMap nl) -> VertexMapRef nl m
VMR forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef forall nl. VertexMap nl
emptyVertexMap

-- | Just like 'vertexForLabel', but holding the mapping in a ref instead
-- of threading it.  Usage is simpler:
--
-- > v <- vertexForLabelRef g m lbl
vertexForLabelRef :: (MLabeledVertex g, Ord (MVertexLabel g), P.PrimMonad m, R.MonadRef m)
                  => g m
                  -> VertexMapRef (MVertexLabel g) m
                  -> MVertexLabel g
                  -> m Vertex
vertexForLabelRef :: forall (g :: (* -> *) -> *) (m :: * -> *).
(MLabeledVertex g, Ord (MVertexLabel g), PrimMonad m,
 MonadRef m) =>
g m
-> VertexMapRef (MVertexLabel g) m -> MVertexLabel g -> m Vertex
vertexForLabelRef g m
g (VMR Ref m (VertexMap (MVertexLabel g))
ref) MVertexLabel g
lbl = do
  VertexMap (MVertexLabel g)
vm <- forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef Ref m (VertexMap (MVertexLabel g))
ref
  (Vertex
v, VertexMap (MVertexLabel g)
vm') <- forall (g :: (* -> *) -> *) (m :: * -> *).
(MLabeledVertex g, Ord (MVertexLabel g), PrimMonad m,
 MonadRef m) =>
g m
-> VertexMap (MVertexLabel g)
-> MVertexLabel g
-> m (Vertex, VertexMap (MVertexLabel g))
vertexForLabel g m
g VertexMap (MVertexLabel g)
vm MVertexLabel g
lbl
  forall (m :: * -> *) a. MonadRef m => Ref m a -> a -> m ()
R.writeRef Ref m (VertexMap (MVertexLabel g))
ref VertexMap (MVertexLabel g)
vm'
  forall (m :: * -> *) a. Monad m => a -> m a
return Vertex
v