{-# 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 Map nl Vertex -> () -> ()
forall a b. NFData a => a -> b -> b
`DS.deepseq` ()

emptyVertexMap :: VertexMap nl
emptyVertexMap :: VertexMap nl
emptyVertexMap = Map nl Vertex -> VertexMap nl
forall nl. Map nl Vertex -> VertexMap nl
VM Map nl Vertex
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 :: 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 <- MVertexLabel g -> Map (MVertexLabel g) Vertex -> Maybe Vertex
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup MVertexLabel g
lbl Map (MVertexLabel g) Vertex
m = (Vertex, VertexMap (MVertexLabel g))
-> m (Vertex, VertexMap (MVertexLabel g))
forall (m :: * -> *) a. Monad m => a -> m a
return (Vertex
v, VertexMap (MVertexLabel g)
vm)
  | Bool
otherwise = do
    Vertex
v <- g m -> MVertexLabel g -> m Vertex
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' = MVertexLabel g
-> Vertex
-> Map (MVertexLabel g) Vertex
-> Map (MVertexLabel g) Vertex
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
    (Vertex, VertexMap (MVertexLabel g))
-> m (Vertex, VertexMap (MVertexLabel g))
forall (m :: * -> *) a. Monad m => a -> m a
return (Vertex
v, Map (MVertexLabel g) Vertex -> VertexMap (MVertexLabel g)
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 :: nl -> VertexMap nl -> Maybe Vertex
lookupVertexForLabel nl
lbl (VM Map nl Vertex
m) = nl -> Map nl Vertex -> Maybe Vertex
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 :: g -> VertexMap (VertexLabel g)
vertexMapFromGraph = Map (VertexLabel g) Vertex -> VertexMap (VertexLabel g)
forall nl. Map nl Vertex -> VertexMap nl
VM (Map (VertexLabel g) Vertex -> VertexMap (VertexLabel g))
-> (g -> Map (VertexLabel g) Vertex)
-> g
-> VertexMap (VertexLabel g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(VertexLabel g, Vertex)] -> Map (VertexLabel g) Vertex
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VertexLabel g, Vertex)] -> Map (VertexLabel g) Vertex)
-> (g -> [(VertexLabel g, Vertex)])
-> g
-> Map (VertexLabel g) Vertex
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Vertex, VertexLabel g) -> (VertexLabel g, Vertex))
-> [(Vertex, VertexLabel g)] -> [(VertexLabel g, Vertex)]
forall a b. (a -> b) -> [a] -> [b]
map (Vertex, VertexLabel g) -> (VertexLabel g, Vertex)
forall a b. (a, b) -> (b, a)
swap ([(Vertex, VertexLabel g)] -> [(VertexLabel g, Vertex)])
-> (g -> [(Vertex, VertexLabel g)])
-> g
-> [(VertexLabel g, Vertex)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> [(Vertex, VertexLabel g)]
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 :: VertexMapRef nl m -> m (VertexMap nl)
vertexMapFromRef (VMR Ref m (VertexMap nl)
ref) = Ref m (VertexMap nl) -> m (VertexMap nl)
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 :: m (VertexMapRef nl m)
newVertexMapRef = (Ref m (VertexMap nl) -> VertexMapRef nl m)
-> m (Ref m (VertexMap nl)) -> m (VertexMapRef nl m)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM Ref m (VertexMap nl) -> VertexMapRef nl m
forall nl (m :: * -> *). Ref m (VertexMap nl) -> VertexMapRef nl m
VMR (m (Ref m (VertexMap nl)) -> m (VertexMapRef nl m))
-> m (Ref m (VertexMap nl)) -> m (VertexMapRef nl m)
forall a b. (a -> b) -> a -> b
$ VertexMap nl -> m (Ref m (VertexMap nl))
forall (m :: * -> *) a. MonadRef m => a -> m (Ref m a)
R.newRef VertexMap nl
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 :: 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 <- Ref m (VertexMap (MVertexLabel g))
-> m (VertexMap (MVertexLabel g))
forall (m :: * -> *) a. MonadRef m => Ref m a -> m a
R.readRef Ref m (VertexMap (MVertexLabel g))
ref
  (Vertex
v, VertexMap (MVertexLabel g)
vm') <- g m
-> VertexMap (MVertexLabel g)
-> MVertexLabel g
-> m (Vertex, VertexMap (MVertexLabel g))
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
  Ref m (VertexMap (MVertexLabel g))
-> VertexMap (MVertexLabel g) -> m ()
forall (m :: * -> *) a. MonadRef m => Ref m a -> a -> m ()
R.writeRef Ref m (VertexMap (MVertexLabel g))
ref VertexMap (MVertexLabel g)
vm'
  Vertex -> m Vertex
forall (m :: * -> *) a. Monad m => a -> m a
return Vertex
v