{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DeriveDataTypeable  #-}
{-# LANGUAGE ScopedTypeVariables #-}

#if __GLASGOW_HASKELL__ >= 701
{-# LANGUAGE Trustworthy        #-}
#endif

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Concurrent.STM.TBQueue
-- Copyright   :  (c) The University of Glasgow 2012
-- License     :  BSD-style (see the file libraries/base/LICENSE)
--
-- Maintainer  :  libraries@haskell.org
-- Stability   :  experimental
-- Portability :  non-portable (requires STM)
--
-- 'TBQueue' is a bounded version of 'TQueue'. The queue has a maximum
-- capacity set when it is created.  If the queue already contains the
-- maximum number of elements, then 'writeTBQueue' retries until an
-- element is removed from the queue.
--
-- The implementation is based on an array to obtain /O(1)/
-- enqueue and dequeue operations.
--
-- @since 2.4
-----------------------------------------------------------------------------

module Control.Concurrent.STM.TBQueue (
    -- * TBQueue
    TBQueue,
    newTBQueue,
    newTBQueueIO,
    readTBQueue,
    tryReadTBQueue,
    flushTBQueue,
    peekTBQueue,
    tryPeekTBQueue,
    writeTBQueue,
    unGetTBQueue,
    lengthTBQueue,
    isEmptyTBQueue,
    isFullTBQueue,
    capacityTBQueue,
  ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative (pure)
#endif
import Data.Array.Base
import Data.Maybe (isJust, isNothing)
import Data.Typeable   (Typeable)
import GHC.Conc
import Numeric.Natural (Natural)
import Prelude         hiding (read)

import Control.Concurrent.STM.TArray

-- | 'TBQueue' is an abstract type representing a bounded FIFO channel.
--
-- @since 2.4
data TBQueue a
   = TBQueue {-# UNPACK #-} !(TVar Int)             -- read index
             {-# UNPACK #-} !(TVar Int)             -- write index
             {-# UNPACK #-} !(TArray Int (Maybe a)) -- elements
             {-# UNPACK #-} !Int                    -- initial capacity
  deriving Typeable

instance Eq (TBQueue a) where
  -- each `TBQueue` has its own `TVar`s, so it's sufficient to compare the first one
  TBQueue TVar Int
a TVar Int
_ TArray Int (Maybe a)
_ Int
_ == :: TBQueue a -> TBQueue a -> Bool
== TBQueue TVar Int
b TVar Int
_ TArray Int (Maybe a)
_ Int
_ = TVar Int
a TVar Int -> TVar Int -> Bool
forall a. Eq a => a -> a -> Bool
== TVar Int
b

-- incMod x cap == (x + 1) `mod` cap
incMod :: Int -> Int -> Int
incMod :: Int -> Int -> Int
incMod Int
x Int
cap = let y :: Int
y = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 in if Int
y Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
cap then Int
0 else Int
y

-- decMod x cap = (x - 1) `mod` cap
decMod :: Int -> Int -> Int
decMod :: Int -> Int -> Int
decMod Int
x Int
cap = if Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
cap Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 else Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-- | Builds and returns a new instance of 'TBQueue'.
newTBQueue :: Natural   -- ^ maximum number of elements the queue can hold
           -> STM (TBQueue a)
newTBQueue :: forall a. Natural -> STM (TBQueue a)
newTBQueue Natural
cap
  | Natural
cap Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
<= Natural
0 = [Char] -> STM (TBQueue a)
forall a. HasCallStack => [Char] -> a
error [Char]
"capacity has to be greater than 0"
  | Natural
cap Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
forall a. Bounded a => a
maxBound :: Int) = [Char] -> STM (TBQueue a)
forall a. HasCallStack => [Char] -> a
error [Char]
"capacity is too big"
  | Bool
otherwise = do
      TVar Int
rindex <- Int -> STM (TVar Int)
forall a. a -> STM (TVar a)
newTVar Int
0
      TVar Int
windex <- Int -> STM (TVar Int)
forall a. a -> STM (TVar a)
newTVar Int
0
      TArray Int (Maybe a)
elements <- (Int, Int) -> Maybe a -> STM (TArray Int (Maybe a))
forall i. Ix i => (i, i) -> Maybe a -> STM (TArray i (Maybe a))
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, Int
cap' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Maybe a
forall a. Maybe a
Nothing
      TBQueue a -> STM (TBQueue a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TVar Int -> TVar Int -> TArray Int (Maybe a) -> Int -> TBQueue a
forall a.
TVar Int -> TVar Int -> TArray Int (Maybe a) -> Int -> TBQueue a
TBQueue TVar Int
rindex TVar Int
windex TArray Int (Maybe a)
elements Int
cap')
 where
  cap' :: Int
cap' = Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
cap

-- | @IO@ version of 'newTBQueue'.  This is useful for creating top-level
-- 'TBQueue's using 'System.IO.Unsafe.unsafePerformIO', because using
-- 'atomically' inside 'System.IO.Unsafe.unsafePerformIO' isn't
-- possible.
newTBQueueIO :: Natural -> IO (TBQueue a)
newTBQueueIO :: forall a. Natural -> IO (TBQueue a)
newTBQueueIO Natural
cap
  | Natural
cap Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
<= Natural
0 = [Char] -> IO (TBQueue a)
forall a. HasCallStack => [Char] -> a
error [Char]
"capacity has to be greater than 0"
  | Natural
cap Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
forall a. Bounded a => a
maxBound :: Int) = [Char] -> IO (TBQueue a)
forall a. HasCallStack => [Char] -> a
error [Char]
"capacity is too big"
  | Bool
otherwise = do
      TVar Int
rindex <- Int -> IO (TVar Int)
forall a. a -> IO (TVar a)
newTVarIO Int
0
      TVar Int
windex <- Int -> IO (TVar Int)
forall a. a -> IO (TVar a)
newTVarIO Int
0
      TArray Int (Maybe a)
elements <- (Int, Int) -> Maybe a -> IO (TArray Int (Maybe a))
forall i. Ix i => (i, i) -> Maybe a -> IO (TArray i (Maybe a))
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, Int
cap' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Maybe a
forall a. Maybe a
Nothing
      TBQueue a -> IO (TBQueue a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TVar Int -> TVar Int -> TArray Int (Maybe a) -> Int -> TBQueue a
forall a.
TVar Int -> TVar Int -> TArray Int (Maybe a) -> Int -> TBQueue a
TBQueue TVar Int
rindex TVar Int
windex TArray Int (Maybe a)
elements Int
cap')
 where
  cap' :: Int
cap' = Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
cap

-- | Write a value to a 'TBQueue'; retries if the queue is full.
writeTBQueue :: TBQueue a -> a -> STM ()
writeTBQueue :: forall a. TBQueue a -> a -> STM ()
writeTBQueue (TBQueue TVar Int
_ TVar Int
windex TArray Int (Maybe a)
elements Int
cap) a
a = do
  Int
w <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
windex
  Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
w
  case Maybe a
ele of
    Maybe a
Nothing -> TArray Int (Maybe a) -> Int -> Maybe a -> STM ()
forall i. Ix i => TArray i (Maybe a) -> Int -> Maybe a -> STM ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite TArray Int (Maybe a)
elements Int
w (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
    Just a
_ -> STM ()
forall a. STM a
retry
  TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Int
windex (Int -> STM ()) -> Int -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
incMod Int
w Int
cap

-- | Read the next value from the 'TBQueue'; retries if the queue is empty.
readTBQueue :: TBQueue a -> STM a
readTBQueue :: forall a. TBQueue a -> STM a
readTBQueue (TBQueue TVar Int
rindex TVar Int
_ TArray Int (Maybe a)
elements Int
cap) = do
  Int
r <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
rindex
  Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
r
  a
a <- case Maybe a
ele of
        Maybe a
Nothing -> STM a
forall a. STM a
retry
        Just a
a -> do
          TArray Int (Maybe a) -> Int -> Maybe a -> STM ()
forall i. Ix i => TArray i (Maybe a) -> Int -> Maybe a -> STM ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite TArray Int (Maybe a)
elements Int
r Maybe a
forall a. Maybe a
Nothing
          a -> STM a
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
  TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Int
rindex (Int -> STM ()) -> Int -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
incMod Int
r Int
cap
  a -> STM a
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a

-- | A version of 'readTBQueue' which does not retry. Instead it
-- returns @Nothing@ if no value is available.
tryReadTBQueue :: TBQueue a -> STM (Maybe a)
tryReadTBQueue :: forall a. TBQueue a -> STM (Maybe a)
tryReadTBQueue TBQueue a
q = (a -> Maybe a) -> STM a -> STM (Maybe a)
forall a b. (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Maybe a
forall a. a -> Maybe a
Just (TBQueue a -> STM a
forall a. TBQueue a -> STM a
readTBQueue TBQueue a
q) STM (Maybe a) -> STM (Maybe a) -> STM (Maybe a)
forall a. STM a -> STM a -> STM a
`orElse` 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

-- | Efficiently read the entire contents of a 'TBQueue' into a list. This
-- function never retries.
--
-- @since 2.4.5
flushTBQueue :: forall a. TBQueue a -> STM [a]
flushTBQueue :: forall a. TBQueue a -> STM [a]
flushTBQueue (TBQueue TVar Int
_rindex TVar Int
windex TArray Int (Maybe a)
elements Int
cap) = do
  Int
w <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
windex
  Int -> [a] -> STM [a]
go (Int -> Int -> Int
decMod Int
w Int
cap) []
 where
  go :: Int -> [a] -> STM [a]
  go :: Int -> [a] -> STM [a]
go Int
i [a]
acc = do
      Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
i
      case Maybe a
ele of
        Maybe a
Nothing -> [a] -> STM [a]
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
acc
        Just a
a -> do
          TArray Int (Maybe a) -> Int -> Maybe a -> STM ()
forall i. Ix i => TArray i (Maybe a) -> Int -> Maybe a -> STM ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite TArray Int (Maybe a)
elements Int
i Maybe a
forall a. Maybe a
Nothing
          Int -> [a] -> STM [a]
go (Int -> Int -> Int
decMod Int
i Int
cap) (a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc)

-- | Get the next value from the @TBQueue@ without removing it,
-- retrying if the queue is empty.
peekTBQueue :: TBQueue a -> STM a
peekTBQueue :: forall a. TBQueue a -> STM a
peekTBQueue (TBQueue TVar Int
rindex TVar Int
_ TArray Int (Maybe a)
elements Int
_) = do
  Int
r <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
rindex
  Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
r
  case Maybe a
ele of
    Maybe a
Nothing -> STM a
forall a. STM a
retry
    Just a
a -> a -> STM a
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a

-- | A version of 'peekTBQueue' which does not retry. Instead it
-- returns @Nothing@ if no value is available.
tryPeekTBQueue :: TBQueue a -> STM (Maybe a)
tryPeekTBQueue :: forall a. TBQueue a -> STM (Maybe a)
tryPeekTBQueue TBQueue a
q = (a -> Maybe a) -> STM a -> STM (Maybe a)
forall a b. (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Maybe a
forall a. a -> Maybe a
Just (TBQueue a -> STM a
forall a. TBQueue a -> STM a
peekTBQueue TBQueue a
q) STM (Maybe a) -> STM (Maybe a) -> STM (Maybe a)
forall a. STM a -> STM a -> STM a
`orElse` 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

-- | Put a data item back onto a channel, where it will be the next item read.
-- Retries if the queue is full.
unGetTBQueue :: TBQueue a -> a -> STM ()
unGetTBQueue :: forall a. TBQueue a -> a -> STM ()
unGetTBQueue (TBQueue TVar Int
rindex TVar Int
_ TArray Int (Maybe a)
elements Int
cap) a
a = do
  Int
r <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
rindex
  Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
r
  case Maybe a
ele of
    Maybe a
Nothing -> TArray Int (Maybe a) -> Int -> Maybe a -> STM ()
forall i. Ix i => TArray i (Maybe a) -> Int -> Maybe a -> STM ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite TArray Int (Maybe a)
elements Int
r (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
    Just a
_ -> STM ()
forall a. STM a
retry
  TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Int
rindex (Int -> STM ()) -> Int -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
decMod Int
r Int
cap

-- | Return the length of a 'TBQueue'.
--
-- @since 2.5.0.0
lengthTBQueue :: TBQueue a -> STM Natural
lengthTBQueue :: forall a. TBQueue a -> STM Natural
lengthTBQueue (TBQueue TVar Int
rindex TVar Int
windex TArray Int (Maybe a)
elements Int
cap) = do
  Int
r <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
rindex
  Int
w <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
windex
  if Int
w Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r then do
    -- length is 0 or cap
    Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
r
    case Maybe a
ele of
      Maybe a
Nothing -> Natural -> STM Natural
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Natural
0
      Just a
_ -> Natural -> STM Natural
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> STM Natural) -> Natural -> STM Natural
forall a b. (a -> b) -> a -> b
$! Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cap
  else do
    let len' :: Int
len' = Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r
    Natural -> STM Natural
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> STM Natural) -> Natural -> STM Natural
forall a b. (a -> b) -> a -> b
$! Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (if Int
len' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Int
len' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
cap else Int
len')

-- | Returns 'True' if the supplied 'TBQueue' is empty.
isEmptyTBQueue :: TBQueue a -> STM Bool
isEmptyTBQueue :: forall a. TBQueue a -> STM Bool
isEmptyTBQueue (TBQueue TVar Int
rindex TVar Int
windex TArray Int (Maybe a)
elements Int
_) = do
  Int
r <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
rindex
  Int
w <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
windex
  if Int
w Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r then do
    Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
r
    Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> STM Bool) -> Bool -> STM Bool
forall a b. (a -> b) -> a -> b
$! Maybe a -> Bool
forall a. Maybe a -> Bool
isNothing Maybe a
ele
  else
    Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | Returns 'True' if the supplied 'TBQueue' is full.
--
-- @since 2.4.3
isFullTBQueue :: TBQueue a -> STM Bool
isFullTBQueue :: forall a. TBQueue a -> STM Bool
isFullTBQueue (TBQueue TVar Int
rindex TVar Int
windex TArray Int (Maybe a)
elements Int
_) = do
  Int
r <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
rindex
  Int
w <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
windex
  if Int
w Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r then do
    Maybe a
ele <- TArray Int (Maybe a) -> Int -> STM (Maybe a)
forall i. Ix i => TArray i (Maybe a) -> Int -> STM (Maybe a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead TArray Int (Maybe a)
elements Int
r
    Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> STM Bool) -> Bool -> STM Bool
forall a b. (a -> b) -> a -> b
$! Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
ele
  else
    Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | The maximum number of elements the queue can hold.
--
-- @since TODO
capacityTBQueue :: TBQueue a -> Natural
capacityTBQueue :: forall a. TBQueue a -> Natural
capacityTBQueue (TBQueue TVar Int
_ TVar Int
_ TArray Int (Maybe a)
_ Int
cap) = Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cap