module Data.LookupTable (
	LookupTable,
	newLookupTable,
	lookupTable, lookupTableM, cacheInTableM,
	hasLookupTable,
	cachedInTable,
	insertTable, insertTableM, storeInTable, storeInTableM
	) where

import Control.Monad.IO.Class
import Control.Concurrent.MVar
import Data.Map (Map)
import qualified Data.Map as M

-- | k-v table
type LookupTable k v = MVar (Map k v)

newLookupTable :: (Ord k, MonadIO m) => m (LookupTable k v)
newLookupTable :: m (LookupTable k v)
newLookupTable = IO (LookupTable k v) -> m (LookupTable k v)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (LookupTable k v) -> m (LookupTable k v))
-> IO (LookupTable k v) -> m (LookupTable k v)
forall a b. (a -> b) -> a -> b
$ Map k v -> IO (LookupTable k v)
forall a. a -> IO (MVar a)
newMVar Map k v
forall a. Monoid a => a
mempty

-- | Lookup, or insert if not exists
lookupTable :: (Ord k, MonadIO m) => k -> v -> LookupTable k v -> m v
lookupTable :: k -> v -> LookupTable k v -> m v
lookupTable k
key v
value LookupTable k v
tbl = IO v -> m v
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO v -> m v) -> IO v -> m v
forall a b. (a -> b) -> a -> b
$ LookupTable k v -> (Map k v -> IO (Map k v, v)) -> IO v
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar LookupTable k v
tbl ((Map k v -> IO (Map k v, v)) -> IO v)
-> (Map k v -> IO (Map k v, v)) -> IO v
forall a b. (a -> b) -> a -> b
$ \Map k v
tbl' -> case k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
key Map k v
tbl' of
	Just v
value' -> (Map k v, v) -> IO (Map k v, v)
forall (m :: * -> *) a. Monad m => a -> m a
return (Map k v
tbl', v
value')
	Maybe v
Nothing -> (Map k v, v) -> IO (Map k v, v)
forall (m :: * -> *) a. Monad m => a -> m a
return (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
key v
value Map k v
tbl', v
value)

-- | Lookup, or insert if not exists
lookupTableM :: (Ord k, MonadIO m) => k -> m v -> LookupTable k v -> m v
lookupTableM :: k -> m v -> LookupTable k v -> m v
lookupTableM k
key m v
mvalue LookupTable k v
tbl = do
	Maybe v
mv <- k -> LookupTable k v -> m (Maybe v)
forall k (m :: * -> *) v.
(Ord k, MonadIO m) =>
k -> LookupTable k v -> m (Maybe v)
hasLookupTable k
key LookupTable k v
tbl
	case Maybe v
mv of
		Just v
value -> v -> m v
forall (m :: * -> *) a. Monad m => a -> m a
return v
value
		Maybe v
Nothing -> do
			v
value <- m v
mvalue
			k -> v -> LookupTable k v -> m v
forall k (m :: * -> *) v.
(Ord k, MonadIO m) =>
k -> v -> LookupTable k v -> m v
lookupTable k
key v
value LookupTable k v
tbl

-- | @lookupTableM@ with swapped args
cacheInTableM :: (Ord k, MonadIO m) => LookupTable k v -> k -> m v -> m v
cacheInTableM :: LookupTable k v -> k -> m v -> m v
cacheInTableM LookupTable k v
tbl k
key m v
mvalue = k -> m v -> LookupTable k v -> m v
forall k (m :: * -> *) v.
(Ord k, MonadIO m) =>
k -> m v -> LookupTable k v -> m v
lookupTableM k
key m v
mvalue LookupTable k v
tbl

-- | Just check existable
hasLookupTable :: (Ord k, MonadIO m) => k -> LookupTable k v -> m (Maybe v)
hasLookupTable :: k -> LookupTable k v -> m (Maybe v)
hasLookupTable k
key LookupTable k v
tbl = IO (Maybe v) -> m (Maybe v)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe v) -> m (Maybe v)) -> IO (Maybe v) -> m (Maybe v)
forall a b. (a -> b) -> a -> b
$ LookupTable k v -> (Map k v -> IO (Maybe v)) -> IO (Maybe v)
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar LookupTable k v
tbl ((Map k v -> IO (Maybe v)) -> IO (Maybe v))
-> (Map k v -> IO (Maybe v)) -> IO (Maybe v)
forall a b. (a -> b) -> a -> b
$ Maybe v -> IO (Maybe v)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe v -> IO (Maybe v))
-> (Map k v -> Maybe v) -> Map k v -> IO (Maybe v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
key

-- | Make function caching results in @LookupTable@
cachedInTable :: (Ord k, MonadIO m) => LookupTable k v -> (k -> m v) -> k -> m v
cachedInTable :: LookupTable k v -> (k -> m v) -> k -> m v
cachedInTable LookupTable k v
tbl k -> m v
fn k
key = LookupTable k v -> k -> m v -> m v
forall k (m :: * -> *) v.
(Ord k, MonadIO m) =>
LookupTable k v -> k -> m v -> m v
cacheInTableM LookupTable k v
tbl k
key (k -> m v
fn k
key)

-- | Insert value into table and return it
insertTable :: (Ord k, MonadIO m) => k -> v -> LookupTable k v -> m v
insertTable :: k -> v -> LookupTable k v -> m v
insertTable k
key v
value LookupTable k v
tbl = IO v -> m v
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO v -> m v) -> IO v -> m v
forall a b. (a -> b) -> a -> b
$ LookupTable k v -> (Map k v -> IO (Map k v, v)) -> IO v
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar LookupTable k v
tbl ((Map k v -> IO (Map k v, v)) -> IO v)
-> (Map k v -> IO (Map k v, v)) -> IO v
forall a b. (a -> b) -> a -> b
$ \Map k v
tbl' -> (Map k v, v) -> IO (Map k v, v)
forall (m :: * -> *) a. Monad m => a -> m a
return (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
key v
value Map k v
tbl', v
value)

-- | Insert value into table and return it
insertTableM :: (Ord k, MonadIO m) => k -> m v -> LookupTable k v -> m v
insertTableM :: k -> m v -> LookupTable k v -> m v
insertTableM k
key m v
mvalue LookupTable k v
tbl = do
	v
value <- m v
mvalue
	k -> v -> LookupTable k v -> m v
forall k (m :: * -> *) v.
(Ord k, MonadIO m) =>
k -> v -> LookupTable k v -> m v
insertTable k
key v
value LookupTable k v
tbl

-- | @insertTable@ with flipped args
storeInTable :: (Ord k, MonadIO m) => LookupTable k v -> k -> v -> m v
storeInTable :: LookupTable k v -> k -> v -> m v
storeInTable LookupTable k v
tbl k
key v
value = k -> v -> LookupTable k v -> m v
forall k (m :: * -> *) v.
(Ord k, MonadIO m) =>
k -> v -> LookupTable k v -> m v
insertTable k
key v
value LookupTable k v
tbl

-- | @insertTable@ with flipped args
storeInTableM :: (Ord k, MonadIO m) => LookupTable k v -> k -> m v -> m v
storeInTableM :: LookupTable k v -> k -> m v -> m v
storeInTableM LookupTable k v
tbl k
key m v
mvalue = k -> m v -> LookupTable k v -> m v
forall k (m :: * -> *) v.
(Ord k, MonadIO m) =>
k -> m v -> LookupTable k v -> m v
insertTableM k
key m v
mvalue LookupTable k v
tbl