{-# LANGUAGE TupleSections #-}

module Simplex.Messaging.TMap
  ( TMap,
    emptyIO,
    singleton,
    clear,
    Simplex.Messaging.TMap.null,
    Simplex.Messaging.TMap.lookup,
    lookupIO,
    member,
    memberIO,
    insert,
    insertM,
    delete,
    lookupInsert,
    lookupDelete,
    adjust,
    update,
    alter,
    alterF,
    union,
  )
where

import Control.Concurrent.STM
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

type TMap k a = TVar (Map k a)

emptyIO :: IO (TMap k a)
emptyIO :: forall k a. IO (TMap k a)
emptyIO = Map k a -> IO (TVar (Map k a))
forall a. a -> IO (TVar a)
newTVarIO Map k a
forall k a. Map k a
M.empty
{-# INLINE emptyIO #-}

singleton :: k -> a -> STM (TMap k a)
singleton :: forall k a. k -> a -> STM (TMap k a)
singleton k
k a
v = Map k a -> STM (TVar (Map k a))
forall a. a -> STM (TVar a)
newTVar (Map k a -> STM (TVar (Map k a)))
-> Map k a -> STM (TVar (Map k a))
forall a b. (a -> b) -> a -> b
$ k -> a -> Map k a
forall k a. k -> a -> Map k a
M.singleton k
k a
v
{-# INLINE singleton #-}

clear :: TMap k a -> STM ()
clear :: forall k a. TMap k a -> STM ()
clear TMap k a
m = TMap k a -> Map k a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TMap k a
m Map k a
forall k a. Map k a
M.empty
{-# INLINE clear #-}

null :: TMap k a -> STM Bool
null :: forall k a. TMap k a -> STM Bool
null TMap k a
m = Map k a -> Bool
forall k a. Map k a -> Bool
M.null (Map k a -> Bool) -> STM (Map k a) -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMap k a -> STM (Map k a)
forall a. TVar a -> STM a
readTVar TMap k a
m
{-# INLINE null #-}

lookup :: Ord k => k -> TMap k a -> STM (Maybe a)
lookup :: forall k a. Ord k => k -> TMap k a -> STM (Maybe a)
lookup k
k TMap k a
m = k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k (Map k a -> Maybe a) -> STM (Map k a) -> STM (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMap k a -> STM (Map k a)
forall a. TVar a -> STM a
readTVar TMap k a
m
{-# INLINE lookup #-}

lookupIO :: Ord k => k -> TMap k a -> IO (Maybe a)
lookupIO :: forall k a. Ord k => k -> TMap k a -> IO (Maybe a)
lookupIO k
k TMap k a
m = k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k (Map k a -> Maybe a) -> IO (Map k a) -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMap k a -> IO (Map k a)
forall a. TVar a -> IO a
readTVarIO TMap k a
m
{-# INLINE lookupIO #-}

member :: Ord k => k -> TMap k a -> STM Bool
member :: forall k a. Ord k => k -> TMap k a -> STM Bool
member k
k TMap k a
m = k -> Map k a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member k
k (Map k a -> Bool) -> STM (Map k a) -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMap k a -> STM (Map k a)
forall a. TVar a -> STM a
readTVar TMap k a
m
{-# INLINE member #-}

memberIO :: Ord k => k -> TMap k a -> IO Bool
memberIO :: forall k a. Ord k => k -> TMap k a -> IO Bool
memberIO k
k TMap k a
m = k -> Map k a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member k
k (Map k a -> Bool) -> IO (Map k a) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMap k a -> IO (Map k a)
forall a. TVar a -> IO a
readTVarIO TMap k a
m
{-# INLINE memberIO #-}

insert :: Ord k => k -> a -> TMap k a -> STM ()
insert :: forall k a. Ord k => k -> a -> TMap k a -> STM ()
insert k
k a
v TMap k a
m = TMap k a -> (Map k a -> Map k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TMap k a
m ((Map k a -> Map k a) -> STM ()) -> (Map k a -> Map k a) -> STM ()
forall a b. (a -> b) -> a -> b
$ k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k a
v
{-# INLINE insert #-}

insertM :: Ord k => k -> STM a -> TMap k a -> STM ()
insertM :: forall k a. Ord k => k -> STM a -> TMap k a -> STM ()
insertM k
k STM a
f TMap k a
m = TMap k a -> (Map k a -> Map k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TMap k a
m ((Map k a -> Map k a) -> STM ())
-> (a -> Map k a -> Map k a) -> a -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (a -> STM ()) -> STM a -> STM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STM a
f
{-# INLINE insertM #-}

delete :: Ord k => k -> TMap k a -> STM ()
delete :: forall k a. Ord k => k -> TMap k a -> STM ()
delete k
k TMap k a
m = TMap k a -> (Map k a -> Map k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TMap k a
m ((Map k a -> Map k a) -> STM ()) -> (Map k a -> Map k a) -> STM ()
forall a b. (a -> b) -> a -> b
$ k -> Map k a -> Map k a
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
k
{-# INLINE delete #-}

lookupInsert :: Ord k => k -> a -> TMap k a -> STM (Maybe a)
lookupInsert :: forall k a. Ord k => k -> a -> TMap k a -> STM (Maybe a)
lookupInsert k
k a
v TMap k a
m = TMap k a -> (Map k a -> (Maybe a, Map k a)) -> STM (Maybe a)
forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar TMap k a
m ((Map k a -> (Maybe a, Map k a)) -> STM (Maybe a))
-> (Map k a -> (Maybe a, Map k a)) -> STM (Maybe a)
forall a b. (a -> b) -> a -> b
$ (Maybe a -> (Maybe a, Maybe a))
-> k -> Map k a -> (Maybe a, Map k a)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
M.alterF (,a -> Maybe a
forall a. a -> Maybe a
Just a
v) k
k
{-# INLINE lookupInsert #-}

lookupDelete :: Ord k => k -> TMap k a -> STM (Maybe a)
lookupDelete :: forall k a. Ord k => k -> TMap k a -> STM (Maybe a)
lookupDelete k
k TMap k a
m = TMap k a -> (Map k a -> (Maybe a, Map k a)) -> STM (Maybe a)
forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar TMap k a
m ((Map k a -> (Maybe a, Map k a)) -> STM (Maybe a))
-> (Map k a -> (Maybe a, Map k a)) -> STM (Maybe a)
forall a b. (a -> b) -> a -> b
$ (Maybe a -> (Maybe a, Maybe a))
-> k -> Map k a -> (Maybe a, Map k a)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
M.alterF (,Maybe a
forall a. Maybe a
Nothing) k
k
{-# INLINE lookupDelete #-}

adjust :: Ord k => (a -> a) -> k -> TMap k a -> STM ()
adjust :: forall k a. Ord k => (a -> a) -> k -> TMap k a -> STM ()
adjust a -> a
f k
k TMap k a
m = TMap k a -> (Map k a -> Map k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TMap k a
m ((Map k a -> Map k a) -> STM ()) -> (Map k a -> Map k a) -> STM ()
forall a b. (a -> b) -> a -> b
$ (a -> a) -> k -> Map k a -> Map k a
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust a -> a
f k
k
{-# INLINE adjust #-}

update :: Ord k => (a -> Maybe a) -> k -> TMap k a -> STM ()
update :: forall k a. Ord k => (a -> Maybe a) -> k -> TMap k a -> STM ()
update a -> Maybe a
f k
k TMap k a
m = TMap k a -> (Map k a -> Map k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TMap k a
m ((Map k a -> Map k a) -> STM ()) -> (Map k a -> Map k a) -> STM ()
forall a b. (a -> b) -> a -> b
$ (a -> Maybe a) -> k -> Map k a -> Map k a
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
M.update a -> Maybe a
f k
k
{-# INLINE update #-}

alter :: Ord k => (Maybe a -> Maybe a) -> k -> TMap k a -> STM ()
alter :: forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> TMap k a -> STM ()
alter Maybe a -> Maybe a
f k
k TMap k a
m = TMap k a -> (Map k a -> Map k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TMap k a
m ((Map k a -> Map k a) -> STM ()) -> (Map k a -> Map k a) -> STM ()
forall a b. (a -> b) -> a -> b
$ (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter Maybe a -> Maybe a
f k
k
{-# INLINE alter #-}

alterF :: Ord k => (Maybe a -> STM (Maybe a)) -> k -> TMap k a -> STM ()
alterF :: forall k a.
Ord k =>
(Maybe a -> STM (Maybe a)) -> k -> TMap k a -> STM ()
alterF Maybe a -> STM (Maybe a)
f k
k TMap k a
m = do
  Map k a
mv <- (Maybe a -> STM (Maybe a)) -> k -> Map k a -> STM (Map k a)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
M.alterF Maybe a -> STM (Maybe a)
f k
k (Map k a -> STM (Map k a)) -> STM (Map k a) -> STM (Map k a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TMap k a -> STM (Map k a)
forall a. TVar a -> STM a
readTVar TMap k a
m
  TMap k a -> Map k a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TMap k a
m (Map k a -> STM ()) -> Map k a -> STM ()
forall a b. (a -> b) -> a -> b
$! Map k a
mv
{-# INLINE alterF #-}

union :: Ord k => Map k a -> TMap k a -> STM ()
union :: forall k a. Ord k => Map k a -> TMap k a -> STM ()
union Map k a
m' TMap k a
m = TMap k a -> (Map k a -> Map k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TMap k a
m ((Map k a -> Map k a) -> STM ()) -> (Map k a -> Map k a) -> STM ()
forall a b. (a -> b) -> a -> b
$ Map k a -> Map k a -> Map k a
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Map k a
m'
{-# INLINE union #-}