module Data.Pool.WaiterQueue
  ( WaiterQueue,
    newQueueIO,
    push,
    pop,
  )
where

import Control.Concurrent.STM

-- | A FIFO queue that supports removing any element from the queue.
--
-- We have a pointer to the head of the list, and a pointer to the
-- final forward pointer in the list.
data WaiterQueue a
  = WaiterQueue
      (TVar (TDList a))
      (TVar (TVar (TDList a)))

-- | Each element has a pointer to the previous element's forward
-- pointer where "previous element" can be a 'TDList' cons cell or the
-- 'WaiterQueue' head pointer.
data TDList a
  = TCons
      (TVar (TVar (TDList a)))
      a
      (TVar (TDList a))
  | TNil

newQueueIO :: IO (WaiterQueue a)
newQueueIO :: IO (WaiterQueue a)
newQueueIO = do
  TVar (TDList a)
emptyVarL <- TDList a -> IO (TVar (TDList a))
forall a. a -> IO (TVar a)
newTVarIO TDList a
forall a. TDList a
TNil
  TVar (TVar (TDList a))
emptyVarR <- TVar (TDList a) -> IO (TVar (TVar (TDList a)))
forall a. a -> IO (TVar a)
newTVarIO TVar (TDList a)
emptyVarL
  WaiterQueue a -> IO (WaiterQueue a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TVar (TDList a) -> TVar (TVar (TDList a)) -> WaiterQueue a
forall a.
TVar (TDList a) -> TVar (TVar (TDList a)) -> WaiterQueue a
WaiterQueue TVar (TDList a)
emptyVarL TVar (TVar (TDList a))
emptyVarR)

removeSelf ::
  -- | 'WaiterQueue's final forward pointer pointer
  TVar (TVar (TDList a)) ->
  -- | Our back pointer
  TVar (TVar (TDList a)) ->
  -- | Our forward pointer
  TVar (TDList a) ->
  STM ()
removeSelf :: TVar (TVar (TDList a))
-> TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
removeSelf TVar (TVar (TDList a))
tv TVar (TVar (TDList a))
prevPP TVar (TDList a)
nextP = do
  TVar (TDList a)
prevP <- TVar (TVar (TDList a)) -> STM (TVar (TDList a))
forall a. TVar a -> STM a
readTVar TVar (TVar (TDList a))
prevPP
  -- If our back pointer points to our forward pointer then we have
  -- already been removed from the queue
  case TVar (TDList a)
prevP TVar (TDList a) -> TVar (TDList a) -> Bool
forall a. Eq a => a -> a -> Bool
== TVar (TDList a)
nextP of
    Bool
True -> () -> STM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Bool
False -> do
      TDList a
next <- TVar (TDList a) -> STM (TDList a)
forall a. TVar a -> STM a
readTVar TVar (TDList a)
nextP
      TVar (TDList a) -> TDList a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TDList a)
prevP TDList a
next
      case TDList a
next of
        TDList a
TNil -> TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TVar (TDList a))
tv TVar (TDList a)
prevP
        TCons TVar (TVar (TDList a))
bp a
_ TVar (TDList a)
_ -> TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TVar (TDList a))
bp TVar (TDList a)
prevP
      TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TVar (TDList a))
prevPP TVar (TDList a)
nextP
{-# INLINE removeSelf #-}

-- | Returns an STM action that removes the pushed element from the
-- queue
push :: WaiterQueue a -> a -> STM (STM ())
push :: WaiterQueue a -> a -> STM (STM ())
push (WaiterQueue TVar (TDList a)
_ TVar (TVar (TDList a))
tv) a
a = do
  TVar (TDList a)
fwdPointer <- TVar (TVar (TDList a)) -> STM (TVar (TDList a))
forall a. TVar a -> STM a
readTVar TVar (TVar (TDList a))
tv
  TVar (TVar (TDList a))
backPointer <- TVar (TDList a) -> STM (TVar (TVar (TDList a)))
forall a. a -> STM (TVar a)
newTVar TVar (TDList a)
fwdPointer
  TVar (TDList a)
emptyVar <- TDList a -> STM (TVar (TDList a))
forall a. a -> STM (TVar a)
newTVar TDList a
forall a. TDList a
TNil
  let cell :: TDList a
cell = TVar (TVar (TDList a)) -> a -> TVar (TDList a) -> TDList a
forall a.
TVar (TVar (TDList a)) -> a -> TVar (TDList a) -> TDList a
TCons TVar (TVar (TDList a))
backPointer a
a TVar (TDList a)
emptyVar
  TVar (TDList a) -> TDList a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TDList a)
fwdPointer TDList a
cell
  TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TVar (TDList a))
tv TVar (TDList a)
emptyVar
  STM () -> STM (STM ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TVar (TVar (TDList a))
-> TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a.
TVar (TVar (TDList a))
-> TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
removeSelf TVar (TVar (TDList a))
tv TVar (TVar (TDList a))
backPointer TVar (TDList a)
emptyVar)
{-# INLINE push #-}

pop :: WaiterQueue a -> STM (Maybe a)
pop :: WaiterQueue a -> STM (Maybe a)
pop (WaiterQueue TVar (TDList a)
hv TVar (TVar (TDList a))
tv) = do
  TDList a
firstElem <- TVar (TDList a) -> STM (TDList a)
forall a. TVar a -> STM a
readTVar TVar (TDList a)
hv
  case TDList a
firstElem of
    TDList a
TNil -> Maybe a -> STM (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
    TCons TVar (TVar (TDList a))
bp a
a TVar (TDList a)
fp -> do
      TDList a
f <- TVar (TDList a) -> STM (TDList a)
forall a. TVar a -> STM a
readTVar TVar (TDList a)
fp
      TVar (TDList a) -> TDList a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TDList a)
hv TDList a
f
      case TDList a
f of
        TDList a
TNil -> TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TVar (TDList a))
tv TVar (TDList a)
hv
        TCons TVar (TVar (TDList a))
fbp a
_ TVar (TDList a)
_ -> TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TVar (TDList a))
fbp TVar (TDList a)
hv
      -- point the back pointer to the forward pointer as a sign that
      -- the cell has been popped (referenced in removeSelf)
      TVar (TVar (TDList a)) -> TVar (TDList a) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (TVar (TDList a))
bp TVar (TDList a)
fp
      Maybe a -> STM (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
{-# INLINE pop #-}