module Control.Concurrent.STM.TBMQueue
  ( TBMQueue
  , newTBMQueue
  , newTBMQueueIO
  , writeTBMQueue
  , isFullTBMQueue
  , flushTBMQueue
  , sizeTBMQueue
  ) where

import Control.Concurrent.STM (STM, retry)
import Control.Concurrent.STM.TQueue (TQueue, flushTQueue, newTQueue, newTQueueIO, writeTQueue)
import Control.Concurrent.STM.TVar (TVar, newTVar, newTVarIO, readTVar, writeTVar)
import Numeric.Natural (Natural)


-- | 'TBMQueue' is an abstract type representing a bounded FIFO channel with a custom measure function.
-- 'TBMQueue' with a measure function 'const 1' should be equivalent to 'TBQueue'.
data TBMQueue a = TBMQueue !(TQueue a)
                           {-# UNPACK #-} !(TVar Natural) -- Current size
                           !Natural -- Max size
                           !(a -> Natural) -- Measure function

-- | Builds and returns a new instance of 'TBMQueue'.
newTBMQueue :: Natural -> (a -> Natural) -> STM (TBMQueue a)
newTBMQueue :: Natural -> (a -> Natural) -> STM (TBMQueue a)
newTBMQueue Natural
maxSize a -> Natural
measure = do
  TQueue a
queue <- STM (TQueue a)
forall a. STM (TQueue a)
newTQueue
  TVar Natural
currentSize <- Natural -> STM (TVar Natural)
forall a. a -> STM (TVar a)
newTVar Natural
0
  TBMQueue a -> STM (TBMQueue a)
forall (m :: * -> *) a. Monad m => a -> m a
return (TBMQueue a -> STM (TBMQueue a)) -> TBMQueue a -> STM (TBMQueue a)
forall a b. (a -> b) -> a -> b
$ TQueue a -> TVar Natural -> Natural -> (a -> Natural) -> TBMQueue a
forall a.
TQueue a -> TVar Natural -> Natural -> (a -> Natural) -> TBMQueue a
TBMQueue TQueue a
queue TVar Natural
currentSize Natural
maxSize a -> Natural
measure

-- | 'IO' version of 'newTBMQueue'.
newTBMQueueIO :: Natural -> (a -> Natural) -> IO (TBMQueue a)
newTBMQueueIO :: Natural -> (a -> Natural) -> IO (TBMQueue a)
newTBMQueueIO Natural
maxSize a -> Natural
measure = do
  TQueue a
queue <- IO (TQueue a)
forall a. IO (TQueue a)
newTQueueIO
  TVar Natural
currentSize <- Natural -> IO (TVar Natural)
forall a. a -> IO (TVar a)
newTVarIO Natural
0
  TBMQueue a -> IO (TBMQueue a)
forall (m :: * -> *) a. Monad m => a -> m a
return (TBMQueue a -> IO (TBMQueue a)) -> TBMQueue a -> IO (TBMQueue a)
forall a b. (a -> b) -> a -> b
$ TQueue a -> TVar Natural -> Natural -> (a -> Natural) -> TBMQueue a
forall a.
TQueue a -> TVar Natural -> Natural -> (a -> Natural) -> TBMQueue a
TBMQueue TQueue a
queue TVar Natural
currentSize Natural
maxSize a -> Natural
measure

-- | Write a value to a 'TBMQueue'; blocks if the queue is full.
writeTBMQueue :: TBMQueue a -> a -> STM ()
writeTBMQueue :: TBMQueue a -> a -> STM ()
writeTBMQueue (TBMQueue TQueue a
q TVar Natural
currentSize Natural
maxSize a -> Natural
measure) a
item = do
  Natural
size <- TVar Natural -> STM Natural
forall a. TVar a -> STM a
readTVar TVar Natural
currentSize
  if Natural
size Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
maxSize then STM ()
forall a. STM a
retry
  else do
    TQueue a -> a -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue a
q a
item
    TVar Natural -> Natural -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Natural
currentSize (a -> Natural
measure a
item Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
size)

-- | Returns 'True' if the supplied 'TBMQueue' is full.
isFullTBMQueue :: TBMQueue a -> STM Bool
isFullTBMQueue :: TBMQueue a -> STM Bool
isFullTBMQueue (TBMQueue TQueue a
_ TVar Natural
currentSize Natural
maxSize a -> Natural
_) = do
  Natural
size <- TVar Natural -> STM Natural
forall a. TVar a -> STM a
readTVar TVar Natural
currentSize
  Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> STM Bool) -> Bool -> STM Bool
forall a b. (a -> b) -> a -> b
$ Natural
size Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
maxSize

-- | Efficiently read the entire contents of a 'TBMQueue' into a list. This
-- function never retries.
flushTBMQueue :: TBMQueue a -> STM [a]
flushTBMQueue :: TBMQueue a -> STM [a]
flushTBMQueue (TBMQueue TQueue a
q TVar Natural
currentSize Natural
_ a -> Natural
_) = do
  [a]
items <- TQueue a -> STM [a]
forall a. TQueue a -> STM [a]
flushTQueue TQueue a
q
  TVar Natural -> Natural -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Natural
currentSize Natural
0
  [a] -> STM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
items

-- | Returns current size of the queue.
sizeTBMQueue :: TBMQueue a -> STM Natural
sizeTBMQueue :: TBMQueue a -> STM Natural
sizeTBMQueue (TBMQueue TQueue a
_ TVar Natural
currentSize Natural
_ a -> Natural
_) = TVar Natural -> STM Natural
forall a. TVar a -> STM a
readTVar TVar Natural
currentSize