module Hasura.GraphQL.Execute.Subscription.TMap
  ( TMap,
    new,
    reset,
    null,
    lookup,
    insert,
    delete,
    toList,
    replace,
    union,
    filterWithKey,
    getMap,
    adjust,
  )
where

import Control.Concurrent.STM
import Data.HashMap.Strict qualified as HashMap
import Hasura.Prelude hiding (lookup, null, toList, union)

-- | A coarse-grained transactional map implemented by simply wrapping a 'HashMap.HashMap' in a 'TVar'.
-- Compared to "StmContainers.Map", this provides much faster iteration over the elements at the
-- cost of significantly increased contention on writes.
newtype TMap k v = TMap {forall k v. TMap k v -> TVar (HashMap k v)
unTMap :: TVar (HashMap.HashMap k v)}

new :: STM (TMap k v)
new :: forall k v. STM (TMap k v)
new = TVar (HashMap k v) -> TMap k v
forall k v. TVar (HashMap k v) -> TMap k v
TMap (TVar (HashMap k v) -> TMap k v)
-> STM (TVar (HashMap k v)) -> STM (TMap k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap k v -> STM (TVar (HashMap k v))
forall a. a -> STM (TVar a)
newTVar HashMap k v
forall k v. HashMap k v
HashMap.empty

reset :: TMap k v -> STM ()
reset :: forall k v. TMap k v -> STM ()
reset = (TVar (HashMap k v) -> HashMap k v -> STM ())
-> HashMap k v -> TVar (HashMap k v) -> STM ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip TVar (HashMap k v) -> HashMap k v -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar HashMap k v
forall k v. HashMap k v
HashMap.empty (TVar (HashMap k v) -> STM ())
-> (TMap k v -> TVar (HashMap k v)) -> TMap k v -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap

null :: TMap k v -> STM Bool
null :: forall k v. TMap k v -> STM Bool
null = (HashMap k v -> Bool) -> STM (HashMap k v) -> STM Bool
forall a b. (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap HashMap k v -> Bool
forall k v. HashMap k v -> Bool
HashMap.null (STM (HashMap k v) -> STM Bool)
-> (TMap k v -> STM (HashMap k v)) -> TMap k v -> STM Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap k v) -> STM (HashMap k v)
forall a. TVar a -> STM a
readTVar (TVar (HashMap k v) -> STM (HashMap k v))
-> (TMap k v -> TVar (HashMap k v))
-> TMap k v
-> STM (HashMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap

lookup :: (Hashable k) => k -> TMap k v -> STM (Maybe v)
lookup :: forall k v. Hashable k => k -> TMap k v -> STM (Maybe v)
lookup k
k = (HashMap k v -> Maybe v) -> STM (HashMap k v) -> STM (Maybe v)
forall a b. (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (k -> HashMap k v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k) (STM (HashMap k v) -> STM (Maybe v))
-> (TMap k v -> STM (HashMap k v)) -> TMap k v -> STM (Maybe v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap k v) -> STM (HashMap k v)
forall a. TVar a -> STM a
readTVar (TVar (HashMap k v) -> STM (HashMap k v))
-> (TMap k v -> TVar (HashMap k v))
-> TMap k v
-> STM (HashMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap

insert :: (Hashable k) => v -> k -> TMap k v -> STM ()
insert :: forall k v. Hashable k => v -> k -> TMap k v -> STM ()
insert !v
v k
k TMap k v
mapTv = TVar (HashMap k v) -> (HashMap k v -> HashMap k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap TMap k v
mapTv) ((HashMap k v -> HashMap k v) -> STM ())
-> (HashMap k v -> HashMap k v) -> STM ()
forall a b. (a -> b) -> a -> b
$ k -> v -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert k
k v
v

delete :: (Hashable k) => k -> TMap k v -> STM ()
delete :: forall k v. Hashable k => k -> TMap k v -> STM ()
delete k
k TMap k v
mapTv = TVar (HashMap k v) -> (HashMap k v -> HashMap k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap TMap k v
mapTv) ((HashMap k v -> HashMap k v) -> STM ())
-> (HashMap k v -> HashMap k v) -> STM ()
forall a b. (a -> b) -> a -> b
$ k -> HashMap k v -> HashMap k v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HashMap.delete k
k

toList :: TMap k v -> STM [(k, v)]
toList :: forall k v. TMap k v -> STM [(k, v)]
toList = (HashMap k v -> [(k, v)]) -> STM (HashMap k v) -> STM [(k, v)]
forall a b. (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap HashMap k v -> [(k, v)]
forall k v. HashMap k v -> [(k, v)]
HashMap.toList (STM (HashMap k v) -> STM [(k, v)])
-> (TMap k v -> STM (HashMap k v)) -> TMap k v -> STM [(k, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (HashMap k v) -> STM (HashMap k v)
forall a. TVar a -> STM a
readTVar (TVar (HashMap k v) -> STM (HashMap k v))
-> (TMap k v -> TVar (HashMap k v))
-> TMap k v
-> STM (HashMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap

filterWithKey :: (k -> v -> Bool) -> TMap k v -> STM ()
filterWithKey :: forall k v. (k -> v -> Bool) -> TMap k v -> STM ()
filterWithKey k -> v -> Bool
f TMap k v
mapTV = TVar (HashMap k v) -> (HashMap k v -> HashMap k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap TMap k v
mapTV) ((HashMap k v -> HashMap k v) -> STM ())
-> (HashMap k v -> HashMap k v) -> STM ()
forall a b. (a -> b) -> a -> b
$ (k -> v -> Bool) -> HashMap k v -> HashMap k v
forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
HashMap.filterWithKey k -> v -> Bool
f

replace :: TMap k v -> HashMap.HashMap k v -> STM ()
replace :: forall k v. TMap k v -> HashMap k v -> STM ()
replace TMap k v
mapTV HashMap k v
v = STM (HashMap k v) -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM (HashMap k v) -> STM ()) -> STM (HashMap k v) -> STM ()
forall a b. (a -> b) -> a -> b
$ TVar (HashMap k v) -> HashMap k v -> STM (HashMap k v)
forall a. TVar a -> a -> STM a
swapTVar (TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap TMap k v
mapTV) HashMap k v
v

union :: (Hashable k) => TMap k v -> TMap k v -> STM (TMap k v)
union :: forall k v. Hashable k => TMap k v -> TMap k v -> STM (TMap k v)
union TMap k v
mapA TMap k v
mapB = do
  HashMap k v
l <- TVar (HashMap k v) -> STM (HashMap k v)
forall a. TVar a -> STM a
readTVar (TVar (HashMap k v) -> STM (HashMap k v))
-> TVar (HashMap k v) -> STM (HashMap k v)
forall a b. (a -> b) -> a -> b
$ TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap TMap k v
mapA
  HashMap k v
r <- TVar (HashMap k v) -> STM (HashMap k v)
forall a. TVar a -> STM a
readTVar (TVar (HashMap k v) -> STM (HashMap k v))
-> TVar (HashMap k v) -> STM (HashMap k v)
forall a b. (a -> b) -> a -> b
$ TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap TMap k v
mapB
  TVar (HashMap k v) -> TMap k v
forall k v. TVar (HashMap k v) -> TMap k v
TMap (TVar (HashMap k v) -> TMap k v)
-> STM (TVar (HashMap k v)) -> STM (TMap k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap k v -> STM (TVar (HashMap k v))
forall a. a -> STM (TVar a)
newTVar (HashMap k v -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
HashMap k v -> HashMap k v -> HashMap k v
HashMap.union HashMap k v
l HashMap k v
r)

getMap :: TMap k v -> STM (HashMap.HashMap k v)
getMap :: forall k v. TMap k v -> STM (HashMap k v)
getMap = TVar (HashMap k v) -> STM (HashMap k v)
forall a. TVar a -> STM a
readTVar (TVar (HashMap k v) -> STM (HashMap k v))
-> (TMap k v -> TVar (HashMap k v))
-> TMap k v
-> STM (HashMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap

adjust :: (Hashable k) => (v -> v) -> k -> TMap k v -> STM ()
adjust :: forall k v. Hashable k => (v -> v) -> k -> TMap k v -> STM ()
adjust v -> v
f k
k TMap k v
mapTV = TVar (HashMap k v) -> (HashMap k v -> HashMap k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (TMap k v -> TVar (HashMap k v)
forall k v. TMap k v -> TVar (HashMap k v)
unTMap TMap k v
mapTV) ((HashMap k v -> HashMap k v) -> STM ())
-> (HashMap k v -> HashMap k v) -> STM ()
forall a b. (a -> b) -> a -> b
$ (v -> v) -> k -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
HashMap.adjust v -> v
f k
k