{-# LANGUAGE StrictData #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Galley.Queue
  ( Queue,
    new,
    tryPush,
    pop,
    len,
    interpretQueue,
  )
where

import Control.Concurrent.STM qualified as Stm
import Galley.Effects.Queue qualified as E
import Imports
import Numeric.Natural (Natural)
import Polysemy

data Queue a = Queue
  { forall a. Queue a -> TVar Word
_len :: Stm.TVar Word,
    forall a. Queue a -> TBQueue a
_queue :: Stm.TBQueue a
  }

new :: (MonadIO m) => Natural -> m (Queue a)
new :: forall (m :: * -> *) a. MonadIO m => Natural -> m (Queue a)
new Natural
n = IO (Queue a) -> m (Queue a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Queue a) -> m (Queue a)) -> IO (Queue a) -> m (Queue a)
forall a b. (a -> b) -> a -> b
$ TVar Word -> TBQueue a -> Queue a
forall a. TVar Word -> TBQueue a -> Queue a
Queue (TVar Word -> TBQueue a -> Queue a)
-> IO (TVar Word) -> IO (TBQueue a -> Queue a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> IO (TVar Word)
forall a. a -> IO (TVar a)
Stm.newTVarIO Word
0 IO (TBQueue a -> Queue a) -> IO (TBQueue a) -> IO (Queue a)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Natural -> IO (TBQueue a)
forall a. Natural -> IO (TBQueue a)
Stm.newTBQueueIO Natural
n

tryPush :: (MonadIO m) => Queue a -> a -> m Bool
tryPush :: forall (m :: * -> *) a. MonadIO m => Queue a -> a -> m Bool
tryPush Queue a
q a
a = IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> (STM Bool -> IO Bool) -> STM Bool -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM Bool -> IO Bool
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Bool -> m Bool) -> STM Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ do
  Bool
isFull <- TBQueue a -> STM Bool
forall a. TBQueue a -> STM Bool
Stm.isFullTBQueue (Queue a -> TBQueue a
forall a. Queue a -> TBQueue a
_queue Queue a
q)
  Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
isFull (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$ do
    TVar Word -> (Word -> Word) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
Stm.modifyTVar' (Queue a -> TVar Word
forall a. Queue a -> TVar Word
_len Queue a
q) Word -> Word
forall a. Enum a => a -> a
succ
    TBQueue a -> a -> STM ()
forall a. TBQueue a -> a -> STM ()
Stm.writeTBQueue (Queue a -> TBQueue a
forall a. Queue a -> TBQueue a
_queue Queue a
q) a
a
  Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> Bool
not Bool
isFull)

pop :: (MonadIO m) => Queue a -> m a
pop :: forall (m :: * -> *) a. MonadIO m => Queue a -> m a
pop Queue a
q = IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> (STM a -> IO a) -> STM a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM a -> IO a
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM a -> m a) -> STM a -> m a
forall a b. (a -> b) -> a -> b
$ do
  TVar Word -> (Word -> Word) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
Stm.modifyTVar' (Queue a -> TVar Word
forall a. Queue a -> TVar Word
_len Queue a
q) (Word -> Word
forall a. Enum a => a -> a
pred (Word -> Word) -> (Word -> Word) -> Word -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> Word
forall a. Ord a => a -> a -> a
max Word
1)
  TBQueue a -> STM a
forall a. TBQueue a -> STM a
Stm.readTBQueue (Queue a -> TBQueue a
forall a. Queue a -> TBQueue a
_queue Queue a
q)

len :: (MonadIO m) => Queue a -> m Word
len :: forall (m :: * -> *) a. MonadIO m => Queue a -> m Word
len Queue a
q = IO Word -> m Word
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Word -> m Word) -> IO Word -> m Word
forall a b. (a -> b) -> a -> b
$ TVar Word -> IO Word
forall a. TVar a -> IO a
Stm.readTVarIO (Queue a -> TVar Word
forall a. Queue a -> TVar Word
_len Queue a
q)

interpretQueue ::
  (Member (Embed IO) r) =>
  Queue a ->
  Sem (E.Queue a ': r) x ->
  Sem r x
interpretQueue :: forall (r :: EffectRow) a x.
Member (Embed IO) r =>
Queue a -> Sem (Queue a : r) x -> Sem r x
interpretQueue Queue a
q = (forall (rInitial :: EffectRow) x.
 Queue a (Sem rInitial) x -> Sem r x)
-> Sem (Queue a : r) x -> Sem r x
forall (e :: (* -> *) -> * -> *) (r :: EffectRow) a.
FirstOrder e "interpret" =>
(forall (rInitial :: EffectRow) x. e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall (rInitial :: EffectRow) x.
  Queue a (Sem rInitial) x -> Sem r x)
 -> Sem (Queue a : r) x -> Sem r x)
-> (forall (rInitial :: EffectRow) x.
    Queue a (Sem rInitial) x -> Sem r x)
-> Sem (Queue a : r) x
-> Sem r x
forall a b. (a -> b) -> a -> b
$ \case
  E.TryPush a
a -> forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed @IO (IO x -> Sem r x) -> IO x -> Sem r x
forall a b. (a -> b) -> a -> b
$ Queue a -> a -> IO Bool
forall (m :: * -> *) a. MonadIO m => Queue a -> a -> m Bool
tryPush Queue a
q a
a
  Queue a (Sem rInitial) x
E.Pop -> forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed @IO (IO x -> Sem r x) -> IO x -> Sem r x
forall a b. (a -> b) -> a -> b
$ Queue x -> IO x
forall (m :: * -> *) a. MonadIO m => Queue a -> m a
pop Queue a
Queue x
q