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 = liftIO $ newMVar mempty

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

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

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

-- | Just check existable
hasLookupTable :: (Ord k, MonadIO m) => k -> LookupTable k v -> m (Maybe v)
hasLookupTable key tbl = liftIO $ withMVar tbl $ return . M.lookup key

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

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

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

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

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