{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE NoLinearTypes #-}

{- | A closable, concurrent double-ended queue backed by STM, with amortized
O(1) operations. The underlying implementation uses a two-stack design with
separate 'TVar's for the front and rear, reducing STM contention: in the
common case, @pushFront@ and @popBack@ touch disjoint variables and do not
conflict.

Closing semantics follow @stm-chans@ conventions:

  * __Closed + empty__ → read returns @Nothing@ (end-of-stream)
  * __Closed + non-empty__ → read returns @Just a@ (drain remaining)
  * __Open + empty__ → read blocks (@retry@)
  * __Open + non-empty__ → read returns @Just a@
  * __Write to closed__ → silently ignored
-}
module Control.Concurrent.STM.TMDeque (
  -- * The TMDeque type
  TMDeque,

  -- * Construction
  newTMDeque,
  newTMDequeIO,

  -- * Push operations
  pushFrontTMDeque,

  -- * Pop operations (blocking)
  popFrontTMDeque,
  popBackTMDeque,

  -- * Pop operations (non-blocking)
  tryPopFrontTMDeque,
  tryPopBackTMDeque,

  -- * Closing & queries
  closeTMDeque,
  isClosedTMDeque,
  isClosedTMDequeIO,
  isEmptyTMDeque,
  sizeTMDeque,
  countTMDequeIO,
) where

import Control.Concurrent.STM (STM, TVar, modifyTVar', newTVar, newTVarIO, readTVar, retry, writeTVar)
import Control.Concurrent.STM.TVar (readTVarIO)
import Control.Monad (unless)

{- | Reverse a non-empty list and split into head and tail.
Precondition: the input list is non-empty.
-}
unconsReverse :: [a] -> (a, [a])
unconsReverse :: forall a. [a] -> (a, [a])
unconsReverse [a]
xs = case [a] -> [a]
forall a. [a] -> [a]
reverse [a]
xs of
  a
y : [a]
ys -> (a
y, [a]
ys)
  [] -> [Char] -> (a, [a])
forall a. HasCallStack => [Char] -> a
error [Char]
"TMDeque.unconsReverse: impossible – called on empty list"

------------------------------------------------------------------------
-- STM two-stack queue
------------------------------------------------------------------------

-- | A closable, STM-backed double-ended queue with amortized O(1) operations.
data TMDeque a
  = TMDeque
      {-# UNPACK #-} !(TVar Bool) -- closed flag (monotonic: False → True)
      {-# UNPACK #-} !(TVar [a]) -- front (push end)
      {-# UNPACK #-} !(TVar [a]) -- rear (pop end)
      {-# UNPACK #-} !(TVar Int) -- size (maintained for O(1) count)

-- | Create a new empty 'TMDeque'.
newTMDeque :: STM (TMDeque a)
newTMDeque :: forall a. STM (TMDeque a)
newTMDeque = TVar Bool -> TVar [a] -> TVar [a] -> TVar Int -> TMDeque a
forall a.
TVar Bool -> TVar [a] -> TVar [a] -> TVar Int -> TMDeque a
TMDeque (TVar Bool -> TVar [a] -> TVar [a] -> TVar Int -> TMDeque a)
-> STM (TVar Bool)
-> STM (TVar [a] -> TVar [a] -> TVar Int -> TMDeque a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> STM (TVar Bool)
forall a. a -> STM (TVar a)
newTVar Bool
False STM (TVar [a] -> TVar [a] -> TVar Int -> TMDeque a)
-> STM (TVar [a]) -> STM (TVar [a] -> TVar Int -> TMDeque a)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [a] -> STM (TVar [a])
forall a. a -> STM (TVar a)
newTVar [] STM (TVar [a] -> TVar Int -> TMDeque a)
-> STM (TVar [a]) -> STM (TVar Int -> TMDeque a)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [a] -> STM (TVar [a])
forall a. a -> STM (TVar a)
newTVar [] STM (TVar Int -> TMDeque a) -> STM (TVar Int) -> STM (TMDeque a)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> STM (TVar Int)
forall a. a -> STM (TVar a)
newTVar Int
0

-- | IO variant of 'newTMDeque'.
newTMDequeIO :: IO (TMDeque a)
newTMDequeIO :: forall a. IO (TMDeque a)
newTMDequeIO = TVar Bool -> TVar [a] -> TVar [a] -> TVar Int -> TMDeque a
forall a.
TVar Bool -> TVar [a] -> TVar [a] -> TVar Int -> TMDeque a
TMDeque (TVar Bool -> TVar [a] -> TVar [a] -> TVar Int -> TMDeque a)
-> IO (TVar Bool)
-> IO (TVar [a] -> TVar [a] -> TVar Int -> TMDeque a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False IO (TVar [a] -> TVar [a] -> TVar Int -> TMDeque a)
-> IO (TVar [a]) -> IO (TVar [a] -> TVar Int -> TMDeque a)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [a] -> IO (TVar [a])
forall a. a -> IO (TVar a)
newTVarIO [] IO (TVar [a] -> TVar Int -> TMDeque a)
-> IO (TVar [a]) -> IO (TVar Int -> TMDeque a)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [a] -> IO (TVar [a])
forall a. a -> IO (TVar a)
newTVarIO [] IO (TVar Int -> TMDeque a) -> IO (TVar Int) -> IO (TMDeque a)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (TVar Int)
forall a. a -> IO (TVar a)
newTVarIO Int
0

{- | Push an element to the front of the deque.  Silently ignored if the
deque is closed.
-}
pushFrontTMDeque :: TMDeque a -> a -> STM ()
pushFrontTMDeque :: forall a. TMDeque a -> a -> STM ()
pushFrontTMDeque (TMDeque TVar Bool
closedVar TVar [a]
frontVar TVar [a]
_rearVar TVar Int
sizeVar) a
x = do
  closed <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
closedVar
  unless closed do
    modifyTVar' frontVar (x :)
    modifyTVar' sizeVar (+ 1)

{- | Pop an element from the front.  Blocks if the deque is open and empty.
Returns @Nothing@ when the deque is closed and empty (end-of-stream).
-}
popFrontTMDeque :: TMDeque a -> STM (Maybe a)
popFrontTMDeque :: forall a. TMDeque a -> STM (Maybe a)
popFrontTMDeque (TMDeque TVar Bool
closedVar TVar [a]
frontVar TVar [a]
rearVar TVar Int
sizeVar) = do
  f <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
frontVar
  case f of
    a
x : [a]
f' -> do
      TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
frontVar [a]
f'
      TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
      Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
x)
    [] -> do
      r <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
rearVar
      case r of
        a
_ : [a]
_ -> do
          let (a
x, [a]
f') = [a] -> (a, [a])
forall a. [a] -> (a, [a])
unconsReverse [a]
r
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
rearVar []
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
frontVar [a]
f'
          TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
          Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
x)
        [] -> do
          closed <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
closedVar
          if closed
            then pure Nothing
            else retry

{- | Non-blocking pop from the front.

  * @Nothing@         — closed (end-of-stream)
  * @Just Nothing@    — open and empty (would block)
  * @Just (Just a)@   — got an element
-}
tryPopFrontTMDeque :: TMDeque a -> STM (Maybe (Maybe a))
tryPopFrontTMDeque :: forall a. TMDeque a -> STM (Maybe (Maybe a))
tryPopFrontTMDeque (TMDeque TVar Bool
closedVar TVar [a]
frontVar TVar [a]
rearVar TVar Int
sizeVar) = do
  f <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
frontVar
  case f of
    a
x : [a]
f' -> do
      TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
frontVar [a]
f'
      TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
      Maybe (Maybe a) -> STM (Maybe (Maybe a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> Maybe (Maybe a)
forall a. a -> Maybe a
Just (a -> Maybe a
forall a. a -> Maybe a
Just a
x))
    [] -> do
      r <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
rearVar
      case r of
        a
_ : [a]
_ -> do
          let (a
x, [a]
f') = [a] -> (a, [a])
forall a. [a] -> (a, [a])
unconsReverse [a]
r
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
rearVar []
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
frontVar [a]
f'
          TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
          Maybe (Maybe a) -> STM (Maybe (Maybe a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> Maybe (Maybe a)
forall a. a -> Maybe a
Just (a -> Maybe a
forall a. a -> Maybe a
Just a
x))
        [] -> do
          closed <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
closedVar
          if closed
            then pure Nothing
            else pure (Just Nothing)

{- | Pop an element from the back.  Blocks if the deque is open and empty.
Returns @Nothing@ when the deque is closed and empty (end-of-stream).
-}
popBackTMDeque :: TMDeque a -> STM (Maybe a)
popBackTMDeque :: forall a. TMDeque a -> STM (Maybe a)
popBackTMDeque (TMDeque TVar Bool
closedVar TVar [a]
frontVar TVar [a]
rearVar TVar Int
sizeVar) = do
  r <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
rearVar
  case r of
    a
x : [a]
r' -> do
      TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
rearVar [a]
r'
      TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
      Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
x)
    [] -> do
      f <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
frontVar
      case f of
        a
_ : [a]
_ -> do
          let (a
x, [a]
r') = [a] -> (a, [a])
forall a. [a] -> (a, [a])
unconsReverse [a]
f
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
frontVar []
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
rearVar [a]
r'
          TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
          Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
x)
        [] -> do
          closed <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
closedVar
          if closed
            then pure Nothing
            else retry

{- | Non-blocking pop from the back.

  * @Nothing@         — closed (end-of-stream)
  * @Just Nothing@    — open and empty (would block)
  * @Just (Just a)@   — got an element
-}
tryPopBackTMDeque :: TMDeque a -> STM (Maybe (Maybe a))
tryPopBackTMDeque :: forall a. TMDeque a -> STM (Maybe (Maybe a))
tryPopBackTMDeque (TMDeque TVar Bool
closedVar TVar [a]
frontVar TVar [a]
rearVar TVar Int
sizeVar) = do
  r <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
rearVar
  case r of
    a
x : [a]
r' -> do
      TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
rearVar [a]
r'
      TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
      Maybe (Maybe a) -> STM (Maybe (Maybe a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> Maybe (Maybe a)
forall a. a -> Maybe a
Just (a -> Maybe a
forall a. a -> Maybe a
Just a
x))
    [] -> do
      f <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
frontVar
      case f of
        a
_ : [a]
_ -> do
          let (a
x, [a]
r') = [a] -> (a, [a])
forall a. [a] -> (a, [a])
unconsReverse [a]
f
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
frontVar []
          TVar [a] -> [a] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar [a]
rearVar [a]
r'
          TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
          Maybe (Maybe a) -> STM (Maybe (Maybe a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> Maybe (Maybe a)
forall a. a -> Maybe a
Just (a -> Maybe a
forall a. a -> Maybe a
Just a
x))
        [] -> do
          closed <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
closedVar
          if closed
            then pure Nothing
            else pure (Just Nothing)

{- | Close the deque.  After closing, writes are silently ignored and reads
will drain remaining elements before signalling end-of-stream.  Closing
is idempotent.
-}
closeTMDeque :: TMDeque a -> STM ()
closeTMDeque :: forall a. TMDeque a -> STM ()
closeTMDeque (TMDeque TVar Bool
closedVar TVar [a]
_ TVar [a]
_ TVar Int
_) = TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
closedVar Bool
True

-- | Check whether the deque has been closed.
isClosedTMDeque :: TMDeque a -> STM Bool
isClosedTMDeque :: forall a. TMDeque a -> STM Bool
isClosedTMDeque (TMDeque TVar Bool
closedVar TVar [a]
_ TVar [a]
_ TVar Int
_) = TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
closedVar

-- | Check whether the deque has been closed.
isClosedTMDequeIO :: TMDeque a -> IO Bool
isClosedTMDequeIO :: forall a. TMDeque a -> IO Bool
isClosedTMDequeIO (TMDeque TVar Bool
closedVar TVar [a]
_ TVar [a]
_ TVar Int
_) = TVar Bool -> IO Bool
forall a. TVar a -> IO a
readTVarIO TVar Bool
closedVar

-- | Check whether the deque is currently empty.
isEmptyTMDeque :: TMDeque a -> STM Bool
isEmptyTMDeque :: forall a. TMDeque a -> STM Bool
isEmptyTMDeque (TMDeque TVar Bool
_ TVar [a]
frontVar TVar [a]
rearVar TVar Int
_) = do
  f <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
frontVar
  case f of
    a
_ : [a]
_ -> Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
    [] -> do
      r <- TVar [a] -> STM [a]
forall a. TVar a -> STM a
readTVar TVar [a]
rearVar
      pure (null r)

-- | Return the number of elements currently in the deque. O(1).
sizeTMDeque :: TMDeque a -> STM Int
sizeTMDeque :: forall a. TMDeque a -> STM Int
sizeTMDeque (TMDeque TVar Bool
_ TVar [a]
_ TVar [a]
_ TVar Int
sizeVar) = TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
sizeVar

-- | IO variant of 'countTMDeque'. O(1).
countTMDequeIO :: TMDeque a -> IO Int
countTMDequeIO :: forall a. TMDeque a -> IO Int
countTMDequeIO (TMDeque TVar Bool
_ TVar [a]
_ TVar [a]
_ TVar Int
sizeVar) = TVar Int -> IO Int
forall a. TVar a -> IO a
readTVarIO TVar Int
sizeVar