{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Simplex.Messaging.Server.MsgStore.STM where

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Simplex.Messaging.Protocol (RecipientId)
import Simplex.Messaging.Server.MsgStore
import UnliftIO.STM

newtype MsgQueue = MsgQueue {MsgQueue -> TQueue Message
msgQueue :: TQueue Message}

newtype MsgStoreData = MsgStoreData {MsgStoreData -> Map RecipientId MsgQueue
messages :: Map RecipientId MsgQueue}

type STMMsgStore = TVar MsgStoreData

newMsgStore :: STM STMMsgStore
newMsgStore :: STM STMMsgStore
newMsgStore = MsgStoreData -> STM STMMsgStore
forall a. a -> STM (TVar a)
newTVar (MsgStoreData -> STM STMMsgStore)
-> MsgStoreData -> STM STMMsgStore
forall a b. (a -> b) -> a -> b
$ Map RecipientId MsgQueue -> MsgStoreData
MsgStoreData Map RecipientId MsgQueue
forall k a. Map k a
M.empty

instance MonadMsgStore STMMsgStore MsgQueue STM where
  getMsgQueue :: STMMsgStore -> RecipientId -> STM MsgQueue
  getMsgQueue :: STMMsgStore -> RecipientId -> STM MsgQueue
getMsgQueue store :: STMMsgStore
store rId :: RecipientId
rId = do
    Map RecipientId MsgQueue
m <- MsgStoreData -> Map RecipientId MsgQueue
messages (MsgStoreData -> Map RecipientId MsgQueue)
-> STM MsgStoreData -> STM (Map RecipientId MsgQueue)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STMMsgStore -> STM MsgStoreData
forall a. TVar a -> STM a
readTVar STMMsgStore
store
    STM MsgQueue
-> (MsgQueue -> STM MsgQueue) -> Maybe MsgQueue -> STM MsgQueue
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Map RecipientId MsgQueue -> STM MsgQueue
newQ Map RecipientId MsgQueue
m) MsgQueue -> STM MsgQueue
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe MsgQueue -> STM MsgQueue) -> Maybe MsgQueue -> STM MsgQueue
forall a b. (a -> b) -> a -> b
$ RecipientId -> Map RecipientId MsgQueue -> Maybe MsgQueue
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup RecipientId
rId Map RecipientId MsgQueue
m
    where
      newQ :: Map RecipientId MsgQueue -> STM MsgQueue
newQ m' :: Map RecipientId MsgQueue
m' = do
        MsgQueue
q <- TQueue Message -> MsgQueue
MsgQueue (TQueue Message -> MsgQueue)
-> STM (TQueue Message) -> STM MsgQueue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM (TQueue Message)
forall a. STM (TQueue a)
newTQueue
        STMMsgStore -> MsgStoreData -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar STMMsgStore
store (MsgStoreData -> STM ())
-> (Map RecipientId MsgQueue -> MsgStoreData)
-> Map RecipientId MsgQueue
-> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map RecipientId MsgQueue -> MsgStoreData
MsgStoreData (Map RecipientId MsgQueue -> STM ())
-> Map RecipientId MsgQueue -> STM ()
forall a b. (a -> b) -> a -> b
$ RecipientId
-> MsgQueue -> Map RecipientId MsgQueue -> Map RecipientId MsgQueue
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert RecipientId
rId MsgQueue
q Map RecipientId MsgQueue
m'
        MsgQueue -> STM MsgQueue
forall (m :: * -> *) a. Monad m => a -> m a
return MsgQueue
q

  delMsgQueue :: STMMsgStore -> RecipientId -> STM ()
  delMsgQueue :: STMMsgStore -> RecipientId -> STM ()
delMsgQueue store :: STMMsgStore
store rId :: RecipientId
rId =
    STMMsgStore -> (MsgStoreData -> MsgStoreData) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar STMMsgStore
store ((MsgStoreData -> MsgStoreData) -> STM ())
-> (MsgStoreData -> MsgStoreData) -> STM ()
forall a b. (a -> b) -> a -> b
$ Map RecipientId MsgQueue -> MsgStoreData
MsgStoreData (Map RecipientId MsgQueue -> MsgStoreData)
-> (MsgStoreData -> Map RecipientId MsgQueue)
-> MsgStoreData
-> MsgStoreData
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RecipientId -> Map RecipientId MsgQueue -> Map RecipientId MsgQueue
forall k a. Ord k => k -> Map k a -> Map k a
M.delete RecipientId
rId (Map RecipientId MsgQueue -> Map RecipientId MsgQueue)
-> (MsgStoreData -> Map RecipientId MsgQueue)
-> MsgStoreData
-> Map RecipientId MsgQueue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MsgStoreData -> Map RecipientId MsgQueue
messages

instance MonadMsgQueue MsgQueue STM where
  writeMsg :: MsgQueue -> Message -> STM ()
  writeMsg :: MsgQueue -> Message -> STM ()
writeMsg = TQueue Message -> Message -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue (TQueue Message -> Message -> STM ())
-> (MsgQueue -> TQueue Message) -> MsgQueue -> Message -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MsgQueue -> TQueue Message
msgQueue

  tryPeekMsg :: MsgQueue -> STM (Maybe Message)
  tryPeekMsg :: MsgQueue -> STM (Maybe Message)
tryPeekMsg = TQueue Message -> STM (Maybe Message)
forall a. TQueue a -> STM (Maybe a)
tryPeekTQueue (TQueue Message -> STM (Maybe Message))
-> (MsgQueue -> TQueue Message) -> MsgQueue -> STM (Maybe Message)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MsgQueue -> TQueue Message
msgQueue

  peekMsg :: MsgQueue -> STM Message
  peekMsg :: MsgQueue -> STM Message
peekMsg = TQueue Message -> STM Message
forall a. TQueue a -> STM a
peekTQueue (TQueue Message -> STM Message)
-> (MsgQueue -> TQueue Message) -> MsgQueue -> STM Message
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MsgQueue -> TQueue Message
msgQueue

  -- atomic delete (== read) last and peek next message if available
  tryDelPeekMsg :: MsgQueue -> STM (Maybe Message)
  tryDelPeekMsg :: MsgQueue -> STM (Maybe Message)
tryDelPeekMsg (MsgQueue q :: TQueue Message
q) = TQueue Message -> STM (Maybe Message)
forall a. TQueue a -> STM (Maybe a)
tryReadTQueue TQueue Message
q STM (Maybe Message) -> STM (Maybe Message) -> STM (Maybe Message)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TQueue Message -> STM (Maybe Message)
forall a. TQueue a -> STM (Maybe a)
tryPeekTQueue TQueue Message
q