-- | A /closeable/, bounded STM queue.
module Control.Concurrent.TBCQueue(
    TBCQueue,
    newTBCQueue,
    newTBCQueueIO,
    readTBCQueue,
    writeTBCQueue,
    closeTBCQueue,
    isOpenTBCQueue,
    isClosedTBCQueue,
    lengthTBCQueue
) where

import Control.Concurrent.STM
import Control.Monad
import Numeric.Natural

data TBCQueue a = TBCQueue {
    forall a. TBCQueue a -> TBQueue a
q :: !(TBQueue a),
    forall a. TBCQueue a -> TVar Bool
open :: !(TVar Bool)
}

newTBCQueue :: Natural -> STM (TBCQueue a)
newTBCQueue :: forall a. Natural -> STM (TBCQueue a)
newTBCQueue Natural
n = TBQueue a -> TVar Bool -> TBCQueue a
forall a. TBQueue a -> TVar Bool -> TBCQueue a
TBCQueue (TBQueue a -> TVar Bool -> TBCQueue a)
-> STM (TBQueue a) -> STM (TVar Bool -> TBCQueue a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Natural -> STM (TBQueue a)
forall a. Natural -> STM (TBQueue a)
newTBQueue Natural
n STM (TVar Bool -> TBCQueue a)
-> STM (TVar Bool) -> STM (TBCQueue a)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> STM (TVar Bool)
forall a. a -> STM (TVar a)
newTVar Bool
True

-- | @atomically . newTBCQueue@
newTBCQueueIO :: Natural -> IO (TBCQueue a)
newTBCQueueIO :: forall a. Natural -> IO (TBCQueue a)
newTBCQueueIO = STM (TBCQueue a) -> IO (TBCQueue a)
forall a. STM a -> IO a
atomically (STM (TBCQueue a) -> IO (TBCQueue a))
-> (Natural -> STM (TBCQueue a)) -> Natural -> IO (TBCQueue a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> STM (TBCQueue a)
forall a. Natural -> STM (TBCQueue a)
newTBCQueue

-- | Returns @Just value@ until the queue is closed, blocking for the next value.
readTBCQueue :: TBCQueue a -> STM (Maybe a)
readTBCQueue :: forall a. TBCQueue a -> STM (Maybe a)
readTBCQueue TBCQueue a
c = do
    Maybe a
mv <- TBQueue a -> STM (Maybe a)
forall a. TBQueue a -> STM (Maybe a)
tryReadTBQueue TBCQueue a
c.q
    case Maybe a
mv of
        Just a
v -> Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> STM (Maybe a)) -> Maybe a -> STM (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
v
        Maybe a
Nothing -> do
            Bool
stillOpen <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TBCQueue a
c.open
            if Bool
stillOpen
                then STM (Maybe a)
forall a. STM a
retry
                else Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

-- | Writes the given value to the queue if it's still open.
--
-- Returns whether the queue is still open so producers know when consumers no longer care.
writeTBCQueue :: TBCQueue a -> a -> STM Bool
writeTBCQueue :: forall a. TBCQueue a -> a -> STM Bool
writeTBCQueue TBCQueue a
c !a
v = do
    Bool
stillOpen <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TBCQueue a
c.open
    Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
stillOpen (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$ TBQueue a -> a -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBCQueue a
c.q a
v
    Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
stillOpen

-- | Closes the queue so that future writes are no-ops and readers get @Nothing@
-- once all previously-written values are read.
closeTBCQueue :: TBCQueue a -> STM ()
closeTBCQueue :: forall a. TBCQueue a -> STM ()
closeTBCQueue TBCQueue a
c = TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TBCQueue a
c.open Bool
False

isOpenTBCQueue :: TBCQueue a -> STM Bool
isOpenTBCQueue :: forall a. TBCQueue a -> STM Bool
isOpenTBCQueue TBCQueue a
c = TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TBCQueue a
c.open

isClosedTBCQueue :: TBCQueue a -> STM Bool
isClosedTBCQueue :: forall a. TBCQueue a -> STM Bool
isClosedTBCQueue = (Bool -> Bool) -> STM Bool -> STM Bool
forall a b. (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> Bool
not (STM Bool -> STM Bool)
-> (TBCQueue a -> STM Bool) -> TBCQueue a -> STM Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TBCQueue a -> STM Bool
forall a. TBCQueue a -> STM Bool
isOpenTBCQueue

-- | The number of items currently in the queue.
--
-- It is obviously racy to use this value outside the current STM transaction.
lengthTBCQueue :: TBCQueue a -> STM Natural
lengthTBCQueue :: forall a. TBCQueue a -> STM Natural
lengthTBCQueue TBCQueue a
c = TBQueue a -> STM Natural
forall a. TBQueue a -> STM Natural
lengthTBQueue TBCQueue a
c.q