-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Graph.PropertyMap
-- Copyright   :  (C) 2011 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  portable
--
-- Total transient monadic maps, used to track information about vertices
-- and edges in a graph
----------------------------------------------------------------------------

module Data.Graph.PropertyMap
  ( PropertyMap(..)
  , modifyP
  , intPropertyMap
  , propertyMap
  , liftPropertyMap
  ) where

import Control.Monad
import Control.Monad.Trans.Class
import qualified Data.IntMap as IntMap
import qualified Data.Map as Map

data PropertyMap m k v = PropertyMap
  { PropertyMap m k v -> k -> m v
getP :: k -> m v
  , PropertyMap m k v -> k -> v -> m (PropertyMap m k v)
putP :: k -> v -> m (PropertyMap m k v)
  }

modifyP :: Monad m => PropertyMap m k v -> k -> (v -> v) -> m (PropertyMap m k v)
modifyP :: PropertyMap m k v -> k -> (v -> v) -> m (PropertyMap m k v)
modifyP PropertyMap m k v
m k
k v -> v
f = do
  v
a <- PropertyMap m k v -> k -> m v
forall (m :: * -> *) k v. PropertyMap m k v -> k -> m v
getP PropertyMap m k v
m k
k
  PropertyMap m k v -> k -> v -> m (PropertyMap m k v)
forall (m :: * -> *) k v.
PropertyMap m k v -> k -> v -> m (PropertyMap m k v)
putP PropertyMap m k v
m k
k (v -> v
f v
a)

-- A pure IntMap-backed vertex map
intPropertyMap :: Monad m => v -> PropertyMap m Int v
intPropertyMap :: v -> PropertyMap m Int v
intPropertyMap v
v0 = v -> IntMap v -> PropertyMap m Int v
forall (m :: * -> *) a.
Monad m =>
a -> IntMap a -> PropertyMap m Int a
go v
v0 IntMap v
forall a. IntMap a
IntMap.empty where
  go :: a -> IntMap a -> PropertyMap m Int a
go a
v IntMap a
m = PropertyMap :: forall (m :: * -> *) k v.
(k -> m v)
-> (k -> v -> m (PropertyMap m k v)) -> PropertyMap m k v
PropertyMap
    { getP :: Int -> m a
getP = \Int
k -> a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$ a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
v a -> a
forall a. a -> a
id (Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k IntMap a
m)
    , putP :: Int -> a -> m (PropertyMap m Int a)
putP = \Int
k a
v' -> PropertyMap m Int a -> m (PropertyMap m Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (PropertyMap m Int a -> m (PropertyMap m Int a))
-> PropertyMap m Int a -> m (PropertyMap m Int a)
forall a b. (a -> b) -> a -> b
$ a -> IntMap a -> PropertyMap m Int a
go a
v (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k a
v' IntMap a
m)
    }

-- A pure Map-backed vertex map
propertyMap :: (Monad m, Ord k) => v -> PropertyMap m k v
propertyMap :: v -> PropertyMap m k v
propertyMap v
v0 = v -> Map k v -> PropertyMap m k v
forall (m :: * -> *) k a.
(Monad m, Ord k) =>
a -> Map k a -> PropertyMap m k a
go v
v0 Map k v
forall k a. Map k a
Map.empty where
  go :: a -> Map k a -> PropertyMap m k a
go a
v Map k a
m = PropertyMap :: forall (m :: * -> *) k v.
(k -> m v)
-> (k -> v -> m (PropertyMap m k v)) -> PropertyMap m k v
PropertyMap
    { getP :: k -> m a
getP = \k
k -> a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$ a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
v a -> a
forall a. a -> a
id (k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k a
m)
    , putP :: k -> a -> m (PropertyMap m k a)
putP = \k
k a
v' -> PropertyMap m k a -> m (PropertyMap m k a)
forall (m :: * -> *) a. Monad m => a -> m a
return (PropertyMap m k a -> m (PropertyMap m k a))
-> PropertyMap m k a -> m (PropertyMap m k a)
forall a b. (a -> b) -> a -> b
$ a -> Map k a -> PropertyMap m k a
go a
v (k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k a
v' Map k a
m)
    }

liftPropertyMap :: (MonadTrans t, Monad m, Monad (t m)) => PropertyMap m k v -> PropertyMap (t m) k v
liftPropertyMap :: PropertyMap m k v -> PropertyMap (t m) k v
liftPropertyMap (PropertyMap k -> m v
g k -> v -> m (PropertyMap m k v)
p) = (k -> t m v)
-> (k -> v -> t m (PropertyMap (t m) k v)) -> PropertyMap (t m) k v
forall (m :: * -> *) k v.
(k -> m v)
-> (k -> v -> m (PropertyMap m k v)) -> PropertyMap m k v
PropertyMap (m v -> t m v
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m v -> t m v) -> (k -> m v) -> k -> t m v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> m v
g) (\k
k v
v -> PropertyMap m k v -> PropertyMap (t m) k v
forall (t :: (* -> *) -> * -> *) (m :: * -> *) k v.
(MonadTrans t, Monad m, Monad (t m)) =>
PropertyMap m k v -> PropertyMap (t m) k v
liftPropertyMap (PropertyMap m k v -> PropertyMap (t m) k v)
-> t m (PropertyMap m k v) -> t m (PropertyMap (t m) k v)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` m (PropertyMap m k v) -> t m (PropertyMap m k v)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (k -> v -> m (PropertyMap m k v)
p k
k v
v))