module Control.Concurrent.STM.TMapChan.Hash where
import Prelude hiding (lookup)
import Data.Hashable (Hashable)
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import Control.Monad (void, forM_)
import Control.Concurrent.STM (STM)
import Control.Concurrent.STM.TChan (TChan, newTChan, writeTChan, readTChan, tryReadTChan, peekTChan, tryPeekTChan, unGetTChan)
import Control.Concurrent.STM.TVar (TVar, newTVar, readTVar, writeTVar, modifyTVar')
newtype TMapChan k a = TMapChan {runTMapChan :: TVar (HashMap k (TChan a))}
newTMapChan :: STM (TMapChan k a)
newTMapChan = TMapChan <$> newTVar HashMap.empty
insert :: (Eq k, Hashable k) => TMapChan k a -> k -> a -> STM ()
insert t k a = do
c <- getTChan t k
writeTChan c a
insertMany :: (Eq k, Hashable k) => TMapChan k a -> k -> [a] -> STM ()
insertMany t k as = forM_ as (insert t k)
insertFirst :: (Eq k, Hashable k) => TMapChan k a -> k -> a -> STM ()
insertFirst t k a = do
c <- getTChan t k
unGetTChan c a
lookup :: (Eq k, Hashable k) => TMapChan k a -> k -> STM a
lookup t k = do
c <- getTChan t k
readTChan c
tryLookup :: (Eq k, Hashable k) => TMapChan k a -> k -> STM (Maybe a)
tryLookup t k = do
c <- getTChan t k
tryReadTChan c
lookupAll :: (Eq k, Hashable k) => TMapChan k a -> k -> STM [a]
lookupAll t k = do
mNext <- tryLookup t k
case mNext of
Nothing -> pure []
Just next -> (next:) <$> lookupAll t k
observe :: (Eq k, Hashable k) => TMapChan k a -> k -> STM a
observe t k = do
c <- getTChan t k
peekTChan c
tryObserve :: (Eq k, Hashable k) => TMapChan k a -> k -> STM (Maybe a)
tryObserve t k = do
c <- getTChan t k
tryPeekTChan c
observeAll :: (Eq k, Hashable k) => TMapChan k a -> k -> STM [a]
observeAll t k = do
mNext <- tryObserve t k
case mNext of
Nothing -> pure []
Just next -> (:) next <$> observeAll t k
delete :: (Eq k, Hashable k) => TMapChan k a -> k -> STM ()
delete t k = void (tryLookup t k)
deleteAll :: (Eq k, Hashable k) => TMapChan k a -> k -> STM ()
deleteAll t k = void (lookupAll t k)
keys :: TMapChan k a -> STM [k]
keys (TMapChan m) = HashMap.keys <$> readTVar m
getTChan :: (Eq k, Hashable k) => TMapChan k a -> k -> STM (TChan a)
getTChan (TMapChan t) k = do
xs <- readTVar t
case HashMap.lookup k xs of
Nothing -> do
c' <- newTChan
writeTVar t (HashMap.insert k c' xs)
pure c'
Just c' -> pure c'
setTChan :: (Eq k, Hashable k) => TMapChan k a -> k -> TChan a -> STM ()
setTChan (TMapChan t) k c = modifyTVar' t (HashMap.insert k c)
broadcast :: (Eq k, Hashable k) => TMapChan k a -> a -> STM ()
broadcast t@(TMapChan xs) a = do
ks <- HashMap.keys <$> readTVar xs
forM_ ks (\k -> insert t k a)
cloneAt :: (Eq k, Hashable k)
=> TMapChan k a
-> k
-> k
-> STM ()
cloneAt t@(TMapChan xs) kFrom kTo = do
as <- observeAll t kFrom
c <- newTChan
forM_ as (writeTChan c)
modifyTVar' xs (HashMap.insert kTo c)
cloneAll :: (Eq k, Hashable k) => TMapChan k a -> k -> STM ()
cloneAll t@(TMapChan xs) k = do
as <- readTVar xs
c <- newTChan
forM_ (HashMap.keys as) $ \k' -> do
as' <- observeAll t k'
forM_ as' (writeTChan c)
modifyTVar' xs (HashMap.insert k c)
cloneAllUniquely :: ( Eq k, Hashable k
, Eq a, Hashable a
)
=> TMapChan k a
-> k
-> STM ()
cloneAllUniquely t@(TMapChan xs) k = do
as <- readTVar xs
c <- newTChan
as' <- newTVar HashSet.empty
forM_ (HashMap.keys as) $ \k' -> do
as'More <- HashSet.fromList <$> observeAll t k'
modifyTVar' as' (HashSet.union as'More)
as'All <- HashSet.toList <$> readTVar as'
forM_ as'All (writeTChan c)
modifyTVar' xs (HashMap.insert k c)