module Metro.IOHashMap
  ( IOHashMap
  , newIOHashMap
  , insert
  , delete
  , lookup
  , update
  , adjust
  , alter
  , null
  , size
  , member
  , keys
  , elems
  , clear
  , toList

  , insertSTM
  , lookupSTM
  , foldrWithKeySTM
  , deleteSTM
  ) where

import           Data.Hashable
import           Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import           Prelude             hiding (lookup, null)
import           UnliftIO            (MonadIO (..), STM, TVar, atomically,
                                      modifyTVar', newTVarIO, readTVar,
                                      readTVarIO)

newtype IOHashMap a b = IOHashMap (TVar (HashMap a b))

newIOHashMap :: MonadIO m => m (IOHashMap a b)
newIOHashMap :: m (IOHashMap a b)
newIOHashMap = TVar (HashMap a b) -> IOHashMap a b
forall a b. TVar (HashMap a b) -> IOHashMap a b
IOHashMap (TVar (HashMap a b) -> IOHashMap a b)
-> m (TVar (HashMap a b)) -> m (IOHashMap a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap a b -> m (TVar (HashMap a b))
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO HashMap a b
forall k v. HashMap k v
HM.empty

insert :: (Eq a, Hashable a, MonadIO m) => IOHashMap a b -> a -> b -> m ()
insert :: IOHashMap a b -> a -> b -> m ()
insert (IOHashMap h :: TVar (HashMap a b)
h) k :: a
k v :: b
v = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ())
-> ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> m ())
-> (HashMap a b -> HashMap a b) -> m ()
forall a b. (a -> b) -> a -> b
$ a -> b -> HashMap a b -> HashMap a b
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert a
k b
v

delete :: (Eq a, Hashable a, MonadIO m) => IOHashMap a b -> a -> m ()
delete :: IOHashMap a b -> a -> m ()
delete (IOHashMap h :: TVar (HashMap a b)
h) k :: a
k = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ())
-> ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> m ())
-> (HashMap a b -> HashMap a b) -> m ()
forall a b. (a -> b) -> a -> b
$ a -> HashMap a b -> HashMap a b
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HM.delete a
k

lookup :: (Eq a, Hashable a, MonadIO m) => IOHashMap a b -> a -> m (Maybe b)
lookup :: IOHashMap a b -> a -> m (Maybe b)
lookup (IOHashMap h :: TVar (HashMap a b)
h) k :: a
k = a -> HashMap a b -> Maybe b
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup a
k (HashMap a b -> Maybe b) -> m (HashMap a b) -> m (Maybe b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> m (HashMap a b)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (HashMap a b)
h

adjust :: (Eq a, Hashable a, MonadIO m) => IOHashMap a b -> (b -> b) -> a -> m ()
adjust :: IOHashMap a b -> (b -> b) -> a -> m ()
adjust (IOHashMap h :: TVar (HashMap a b)
h) f :: b -> b
f k :: a
k = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ())
-> ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> m ())
-> (HashMap a b -> HashMap a b) -> m ()
forall a b. (a -> b) -> a -> b
$ (b -> b) -> a -> HashMap a b -> HashMap a b
forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
HM.adjust b -> b
f a
k

update :: (Eq a, Hashable a, MonadIO m) => IOHashMap a b -> (b -> Maybe b) -> a -> m ()
update :: IOHashMap a b -> (b -> Maybe b) -> a -> m ()
update (IOHashMap h :: TVar (HashMap a b)
h) f :: b -> Maybe b
f k :: a
k = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ())
-> ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> m ())
-> (HashMap a b -> HashMap a b) -> m ()
forall a b. (a -> b) -> a -> b
$ (b -> Maybe b) -> a -> HashMap a b -> HashMap a b
forall k a.
(Eq k, Hashable k) =>
(a -> Maybe a) -> k -> HashMap k a -> HashMap k a
HM.update b -> Maybe b
f a
k

alter :: (Eq a, Hashable a, MonadIO m) => IOHashMap a b -> (Maybe b -> Maybe b) -> a -> m ()
alter :: IOHashMap a b -> (Maybe b -> Maybe b) -> a -> m ()
alter (IOHashMap h :: TVar (HashMap a b)
h) f :: Maybe b -> Maybe b
f k :: a
k = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ())
-> ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> m ())
-> (HashMap a b -> HashMap a b) -> m ()
forall a b. (a -> b) -> a -> b
$ (Maybe b -> Maybe b) -> a -> HashMap a b -> HashMap a b
forall k v.
(Eq k, Hashable k) =>
(Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
HM.alter Maybe b -> Maybe b
f a
k

null :: MonadIO m => IOHashMap a b -> m Bool
null :: IOHashMap a b -> m Bool
null (IOHashMap h :: TVar (HashMap a b)
h) = HashMap a b -> Bool
forall k v. HashMap k v -> Bool
HM.null (HashMap a b -> Bool) -> m (HashMap a b) -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> m (HashMap a b)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (HashMap a b)
h

size :: MonadIO m => IOHashMap a b -> m Int
size :: IOHashMap a b -> m Int
size (IOHashMap h :: TVar (HashMap a b)
h) = HashMap a b -> Int
forall k v. HashMap k v -> Int
HM.size (HashMap a b -> Int) -> m (HashMap a b) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> m (HashMap a b)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (HashMap a b)
h

member :: (Eq a, Hashable a, MonadIO m) => IOHashMap a b -> a -> m Bool
member :: IOHashMap a b -> a -> m Bool
member (IOHashMap h :: TVar (HashMap a b)
h) k :: a
k = a -> HashMap a b -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
HM.member a
k (HashMap a b -> Bool) -> m (HashMap a b) -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> m (HashMap a b)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (HashMap a b)
h

keys :: MonadIO m => IOHashMap a b -> m [a]
keys :: IOHashMap a b -> m [a]
keys (IOHashMap h :: TVar (HashMap a b)
h) = HashMap a b -> [a]
forall k v. HashMap k v -> [k]
HM.keys (HashMap a b -> [a]) -> m (HashMap a b) -> m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> m (HashMap a b)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (HashMap a b)
h

elems :: MonadIO m => IOHashMap a b -> m [b]
elems :: IOHashMap a b -> m [b]
elems (IOHashMap h :: TVar (HashMap a b)
h) = HashMap a b -> [b]
forall k v. HashMap k v -> [v]
HM.elems (HashMap a b -> [b]) -> m (HashMap a b) -> m [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> m (HashMap a b)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (HashMap a b)
h

clear :: MonadIO m => IOHashMap a b -> m ()
clear :: IOHashMap a b -> m ()
clear (IOHashMap h :: TVar (HashMap a b)
h) = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ())
-> ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> m ())
-> (HashMap a b -> HashMap a b) -> m ()
forall a b. (a -> b) -> a -> b
$ HashMap a b -> HashMap a b -> HashMap a b
forall a b. a -> b -> a
const HashMap a b
forall k v. HashMap k v
HM.empty

toList :: MonadIO m => IOHashMap a b -> m [(a, b)]
toList :: IOHashMap a b -> m [(a, b)]
toList (IOHashMap h :: TVar (HashMap a b)
h) = HashMap a b -> [(a, b)]
forall k v. HashMap k v -> [(k, v)]
HM.toList (HashMap a b -> [(a, b)]) -> m (HashMap a b) -> m [(a, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> m (HashMap a b)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar (HashMap a b)
h

insertSTM :: (Eq a, Hashable a) => IOHashMap a b -> a -> b -> STM ()
insertSTM :: IOHashMap a b -> a -> b -> STM ()
insertSTM (IOHashMap h :: TVar (HashMap a b)
h) k :: a
k v :: b
v = TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b) -> STM ()
forall a b. (a -> b) -> a -> b
$ a -> b -> HashMap a b -> HashMap a b
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert a
k b
v

lookupSTM :: (Eq a, Hashable a) => IOHashMap a b -> a -> STM (Maybe b)
lookupSTM :: IOHashMap a b -> a -> STM (Maybe b)
lookupSTM (IOHashMap h :: TVar (HashMap a b)
h) k :: a
k = a -> HashMap a b -> Maybe b
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup a
k (HashMap a b -> Maybe b) -> STM (HashMap a b) -> STM (Maybe b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> STM (HashMap a b)
forall a. TVar a -> STM a
readTVar TVar (HashMap a b)
h

foldrWithKeySTM :: IOHashMap a b -> (a -> b -> c -> c) -> c -> STM c
foldrWithKeySTM :: IOHashMap a b -> (a -> b -> c -> c) -> c -> STM c
foldrWithKeySTM (IOHashMap h :: TVar (HashMap a b)
h) f :: a -> b -> c -> c
f acc :: c
acc = (a -> b -> c -> c) -> c -> HashMap a b -> c
forall k v a. (k -> v -> a -> a) -> a -> HashMap k v -> a
HM.foldrWithKey a -> b -> c -> c
f c
acc (HashMap a b -> c) -> STM (HashMap a b) -> STM c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (HashMap a b) -> STM (HashMap a b)
forall a. TVar a -> STM a
readTVar TVar (HashMap a b)
h

deleteSTM :: (Eq a, Hashable a) => IOHashMap a b -> a -> STM ()
deleteSTM :: IOHashMap a b -> a -> STM ()
deleteSTM (IOHashMap h :: TVar (HashMap a b)
h) k :: a
k = TVar (HashMap a b) -> (HashMap a b -> HashMap a b) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (HashMap a b)
h ((HashMap a b -> HashMap a b) -> STM ())
-> (HashMap a b -> HashMap a b) -> STM ()
forall a b. (a -> b) -> a -> b
$ a -> HashMap a b -> HashMap a b
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HM.delete a
k