-- | A transactional priority queue, based on a Finger Tree.
module Control.Concurrent.STM.TPQueue
  ( TPQueue ()
  , newTPQueue
  , newTPQueueIO
  , writeTPQueue
  , readTPQueue
  , tryReadTPQueue
  , peekTPQueue
  , tryPeekTPQueue
  , isEmptyTPQueue
  ) where

import           Control.Concurrent.STM.TVar
import           Control.Monad.STM
import           Data.PriorityQueue.FingerTree (PQueue)
import qualified Data.PriorityQueue.FingerTree as PQueue

-- | 'TPQueue' is an unbounded priority queue based on a Finger Tree.
newtype TPQueue k v = TPQueue (TVar (PQueue k v))

mkTPQueue :: Functor f => f (TVar (PQueue k v)) -> f (TPQueue k v)
mkTPQueue :: f (TVar (PQueue k v)) -> f (TPQueue k v)
mkTPQueue = (TVar (PQueue k v) -> TPQueue k v)
-> f (TVar (PQueue k v)) -> f (TPQueue k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TVar (PQueue k v) -> TPQueue k v
forall k v. TVar (PQueue k v) -> TPQueue k v
TPQueue

-- | Build a new 'TPQueue'.
newTPQueue :: Ord k => STM (TPQueue k v)
newTPQueue :: STM (TPQueue k v)
newTPQueue = STM (TVar (PQueue k v)) -> STM (TPQueue k v)
forall (f :: * -> *) k v.
Functor f =>
f (TVar (PQueue k v)) -> f (TPQueue k v)
mkTPQueue (PQueue k v -> STM (TVar (PQueue k v))
forall a. a -> STM (TVar a)
newTVar PQueue k v
forall k v. Ord k => PQueue k v
PQueue.empty)

-- | IO version of 'newTPQueue'. This is useful for creating top-level
-- 'TPQueues' using 'unsafePerformIO', because using 'atomically' inside
-- 'unsafePerformIO' isn't possible.
newTPQueueIO :: Ord k => IO (TPQueue k v)
newTPQueueIO :: IO (TPQueue k v)
newTPQueueIO = IO (TVar (PQueue k v)) -> IO (TPQueue k v)
forall (f :: * -> *) k v.
Functor f =>
f (TVar (PQueue k v)) -> f (TPQueue k v)
mkTPQueue (PQueue k v -> IO (TVar (PQueue k v))
forall a. a -> IO (TVar a)
newTVarIO PQueue k v
forall k v. Ord k => PQueue k v
PQueue.empty)

-- | Write a value to a 'TPQueue'.
writeTPQueue :: Ord k => TPQueue k v -> k -> v -> STM ()
writeTPQueue :: TPQueue k v -> k -> v -> STM ()
writeTPQueue (TPQueue TVar (PQueue k v)
h) k
k v
v = TVar (PQueue k v) -> (PQueue k v -> PQueue k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (PQueue k v)
h (k -> v -> PQueue k v -> PQueue k v
forall k v. Ord k => k -> v -> PQueue k v -> PQueue k v
PQueue.add k
k v
v)

-- | Read the next minimal value from a 'TPQueue'.
readTPQueue :: Ord k => TPQueue k v -> STM v
readTPQueue :: TPQueue k v -> STM v
readTPQueue (TPQueue TVar (PQueue k v)
h) = do
    PQueue k v
xs <- TVar (PQueue k v) -> STM (PQueue k v)
forall a. TVar a -> STM a
readTVar TVar (PQueue k v)
h
    case PQueue k v -> Maybe (v, PQueue k v)
forall k v. Ord k => PQueue k v -> Maybe (v, PQueue k v)
PQueue.minView PQueue k v
xs of
        Just (v
x, PQueue k v
xs') -> TVar (PQueue k v) -> PQueue k v -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (PQueue k v)
h PQueue k v
xs' STM () -> STM v -> STM v
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> v -> STM v
forall (f :: * -> *) a. Applicative f => a -> f a
pure v
x
        Maybe (v, PQueue k v)
Nothing       -> STM v
forall a. STM a
retry

-- | A version of 'readTPQueue' that does not retry, but returns 'Nothing'
-- instead if no value is available.
tryReadTPQueue :: Ord k => TPQueue k v -> STM (Maybe v)
tryReadTPQueue :: TPQueue k v -> STM (Maybe v)
tryReadTPQueue (TPQueue TVar (PQueue k v)
h) = do
    PQueue k v
xs <- TVar (PQueue k v) -> STM (PQueue k v)
forall a. TVar a -> STM a
readTVar TVar (PQueue k v)
h
    case PQueue k v -> Maybe (v, PQueue k v)
forall k v. Ord k => PQueue k v -> Maybe (v, PQueue k v)
PQueue.minView PQueue k v
xs of
        Just (v
x, PQueue k v
xs') -> TVar (PQueue k v) -> PQueue k v -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (PQueue k v)
h PQueue k v
xs' STM () -> STM (Maybe v) -> STM (Maybe v)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe v -> STM (Maybe v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v -> Maybe v
forall a. a -> Maybe a
Just v
x)
        Maybe (v, PQueue k v)
Nothing       -> Maybe v -> STM (Maybe v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe v
forall a. Maybe a
Nothing

-- | Get the next minimal value from a 'TPQueue' without removing it.
peekTPQueue :: Ord k => TPQueue k v -> STM v
peekTPQueue :: TPQueue k v -> STM v
peekTPQueue (TPQueue TVar (PQueue k v)
h) = do
    PQueue k v
xs <- TVar (PQueue k v) -> STM (PQueue k v)
forall a. TVar a -> STM a
readTVar TVar (PQueue k v)
h
    case PQueue k v -> Maybe (v, PQueue k v)
forall k v. Ord k => PQueue k v -> Maybe (v, PQueue k v)
PQueue.minView PQueue k v
xs of
        Just (v
x, PQueue k v
_) -> v -> STM v
forall (f :: * -> *) a. Applicative f => a -> f a
pure v
x
        Maybe (v, PQueue k v)
Nothing     -> STM v
forall a. STM a
retry

-- | A version of 'peekTPQueue' that does not retry, but returns 'Nothing'
-- instead if no value is available.
tryPeekTPQueue :: Ord k => TPQueue k v -> STM (Maybe v)
tryPeekTPQueue :: TPQueue k v -> STM (Maybe v)
tryPeekTPQueue (TPQueue TVar (PQueue k v)
h) = do
    PQueue k v
xs <- TVar (PQueue k v) -> STM (PQueue k v)
forall a. TVar a -> STM a
readTVar TVar (PQueue k v)
h
    case PQueue k v -> Maybe (v, PQueue k v)
forall k v. Ord k => PQueue k v -> Maybe (v, PQueue k v)
PQueue.minView PQueue k v
xs of
        Just (v
x, PQueue k v
_) -> Maybe v -> STM (Maybe v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v -> Maybe v
forall a. a -> Maybe a
Just v
x)
        Maybe (v, PQueue k v)
Nothing     -> Maybe v -> STM (Maybe v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe v
forall a. Maybe a
Nothing

-- | Returns 'True' if the 'TPQueue' is empty.
isEmptyTPQueue :: Ord k => TPQueue k v -> STM Bool
isEmptyTPQueue :: TPQueue k v -> STM Bool
isEmptyTPQueue (TPQueue TVar (PQueue k v)
h) = (PQueue k v -> Bool) -> STM (PQueue k v) -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PQueue k v -> Bool
forall k v. Ord k => PQueue k v -> Bool
PQueue.null (TVar (PQueue k v) -> STM (PQueue k v)
forall a. TVar a -> STM a
readTVar TVar (PQueue k v)
h)