module Data.PriorityQueue
( Enqueue(..)
, Dequeue(..)
, DequeueWhere(..)
, PeekQueue(..)
, QueueSize(..)
, PQ
, emptyPQ
, mkPriorityQueue
, mkDefaultPriorityQueue
, PriorityQueue
, newPriorityQueue
, newPriorityQueueBy
) where
import Data.Queue.Classes
import Data.StateRef
import Data.Ord.ReOrd
import qualified Data.Map as M
import Data.List
data PQ a = forall p. Ord p =>
PQ { priorityFunc :: a -> p
, queue :: M.Map p [a]
}
emptyPQ :: Ord p => (a -> p) -> PQ a
emptyPQ f = PQ f M.empty
data PriorityQueue m a =
forall sr. ( ModifyRef sr m (PQ a)
) => PriorityQueue sr
mkPriorityQueue :: ModifyRef sr m (PQ a) => sr -> PriorityQueue m a
mkPriorityQueue = PriorityQueue
mkDefaultPriorityQueue ::
( DefaultStateRef sr m (PQ a)
, ModifyRef sr m (PQ a)
) => sr -> PriorityQueue m a
mkDefaultPriorityQueue = PriorityQueue
newPriorityQueue ::
( DefaultStateRef sr m1 (PQ a)
, ModifyRef sr m1 (PQ a)
, NewRef sr m (PQ a)
, Ord p
) => (a -> p) -> m (PriorityQueue m1 a)
newPriorityQueue f = do
pq <- newRef (emptyPQ f)
return (mkDefaultPriorityQueue pq)
newPriorityQueueBy ::
( DefaultStateRef sr m1 (PQ a)
, ModifyRef sr m1 (PQ a)
, NewRef sr m (PQ a)
) => (a -> a -> Ordering) -> m (PriorityQueue m1 a)
newPriorityQueueBy cmp = newPriorityQueue (ReOrd cmp)
instance Monad m => Enqueue (PriorityQueue m a) m a where
enqueue (PriorityQueue pqRef) x = modifyRef pqRef ins
where ins (PQ f pq) = PQ f (M.insertWith (flip (++)) (f x) [x] pq)
instance Monad m => Dequeue (PriorityQueue m a) m a where
dequeue q@(PriorityQueue pqRef) = atomicModifyRef pqRef dq
where
dq orig@(PQ f pq) = case minViewWithKey pq of
Nothing -> (orig, Nothing)
Just ((k,[]), pq') ->
dq (PQ f pq')
Just ((k,[i]), pq') -> (PQ f pq', Just i)
Just ((k,i:is), pq') -> (PQ f (M.insert k is pq'), Just i)
dequeueBatch q@(PriorityQueue pqRef) = atomicModifyRef pqRef dq
where
dq orig@(PQ f pq) = case M.minView pq of
Nothing -> (orig, [])
Just ([], pq') ->
dq (PQ f pq')
Just (xs, pq') -> (PQ f pq', xs)
instance Monad m => DequeueWhere (PriorityQueue m a) m a where
dequeueWhere (PriorityQueue pqRef) p = atomicModifyRef pqRef dq
where
extractFirstWhere :: (a -> Bool) -> [a] -> (a, [a])
extractFirstWhere p (x:xs)
| p x = (x, xs)
| otherwise = case extractFirstWhere p xs of
(y, rest) -> (y, x:rest)
dq orig@(PQ f pq) = case partition (any p.snd) (M.toAscList pq) of
([], _) ->
(orig, Nothing)
((k, firstMatch): otherMatches, rest) -> case extractFirstWhere p firstMatch of
(thing, otherThings) ->
(PQ f (M.fromList ((k, otherThings) : otherMatches ++ rest)), Just thing)
instance Monad m => PeekQueue (PriorityQueue m a) m a where
peekQueue (PriorityQueue pqRef) = do
PQ f pq <- readRef pqRef
return [v | (k, vs) <- M.toAscList pq, v <- vs]
instance Monad m => QueueSize (PriorityQueue m a) m where
queueSize (PriorityQueue pqRef) = do
PQ f pq <- readRef pqRef
return (M.size pq)
minViewWithKey :: M.Map k a -> Maybe ((k, a), M.Map k a)
#ifdef NoMinViewWithKey
minViewWithKey m = if M.null m
then fail "empty map"
else return (M.deleteFindMin m)
#else
minViewWithKey = M.minViewWithKey
#endif