module Data.PriorityQueue
( Enqueue(..)
, Dequeue(..)
, DequeueWhere(..)
, PeekQueue(..)
, QueueSize(..)
, PQ
, emptyPQ
, mkPriorityQueue
, mkDefaultPriorityQueue
, PriorityQueue
, newPriorityQueue
, newPriorityQueueBy
) where
import Data.Queue.Classes
import Control.Arrow ((&&&), (***))
import Data.StateRef
import Data.Ord.ReOrd
import qualified Data.Map as M
import qualified Data.Sequence as Seq
import Data.Sequence (Seq, viewl, ViewL(..), (><), (<|), singleton, fromList)
import Data.List as List
import Data.Foldable as Foldable
data PQ a = forall p. Ord p =>
PQ { priorityFunc :: a -> p
, queue :: M.Map p (Seq 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 :: Ref m (PQ a) -> PriorityQueue m a
mkDefaultPriorityQueue = PriorityQueue
newPriorityQueue ::
( Monad m
, HasRef m1
, NewRef (Ref m1 (PQ a)) m (PQ a)
, Ord p
) => (a -> p) -> m (PriorityQueue m1 a)
newPriorityQueue f = do
pq <- newReference (emptyPQ f)
return (mkDefaultPriorityQueue pq)
newPriorityQueueBy ::
( Monad m
, HasRef m1
, NewRef (Ref m1 (PQ a)) 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 = modifyReference pqRef $ \(PQ f pq) ->
PQ f (M.insertWith (flip (><)) (f x) (singleton x) pq)
enqueueBatch (PriorityQueue pqRef) xs = modifyReference pqRef $ \(PQ f pq) ->
let prioritized = map (f &&& id) xs
grouped = groupBy ((==) `on` fst) prioritized
batches = map ((head *** fromList) . unzip) grouped
newItems = M.fromListWith (flip (><)) batches
in PQ f (M.unionWith (><) pq newItems)
instance Monad m => Dequeue (PriorityQueue m a) m a where
dequeue (PriorityQueue pqRef) = atomicModifyReference pqRef $ \orig@(PQ f pq) ->
case minViewWithKey pq of
Nothing -> (orig, Nothing)
Just ((k,vs), pq') -> case viewl vs of
EmptyL -> error "dequeue(PriorityQueue): internal inconsistency!"
i :< is
| Seq.null is -> (PQ f pq', Just i)
| otherwise -> (PQ f (M.insert k is pq'), Just i)
dequeueBatch (PriorityQueue pqRef) = atomicModifyReference pqRef $ \orig@(PQ f pq) ->
case M.minView pq of
Nothing -> (orig, [])
Just (xs, pq')
| Seq.null xs -> error "dequeueBatch(PriorityQueue): internal inconsistency!"
| otherwise -> (PQ f pq', toList xs)
instance Monad m => DequeueWhere (PriorityQueue m a) m a where
dequeueWhere (PriorityQueue pqRef) p = atomicModifyReference pqRef $ \orig@(PQ f pq) ->
case List.break (Foldable.any p.snd) (M.toAscList pq) of
(_, []) -> (orig, Nothing)
(nonMatches, (k, firstMatch): rest) -> case extractFirstWhere p firstMatch of
(thing, otherThings)
| Seq.null otherThings ->
(PQ f (M.fromAscList (nonMatches ++ rest)), Just thing)
| otherwise ->
(PQ f (M.fromAscList (nonMatches ++ (k, otherThings) : rest)), Just thing)
instance Monad m => PeekQueue (PriorityQueue m a) m a where
peekQueue (PriorityQueue pqRef) = do
PQ f pq <- readReference pqRef
return [v | (k, vs) <- M.toAscList pq, v <- toList vs]
instance Monad m => QueueSize (PriorityQueue m a) m where
queueSize (PriorityQueue pqRef) = do
PQ f pq <- readReference pqRef
return (M.fold (\xs t -> Seq.length xs + t) 0 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
breakl :: (a -> Bool) -> Seq a -> (Seq a, Seq a)
#ifdef NoBreakL
breakl p xs = case viewl xs of
EmptyL -> (xs, xs)
x :< xs'
| p x -> (Seq.empty, xs)
| otherwise -> let (ys, zs) = breakl p xs' in (x <| ys, zs)
#else
breakl = Seq.breakl
#endif
on :: (b -> b -> c) -> (a -> b) -> (a -> a -> c)
op `on` f = \x y -> f x `op` f y
extractFirstWhere :: (a -> Bool) -> Seq a -> (a, Seq a)
extractFirstWhere p xs = case breakl p xs of
(noMatch, rest) -> case viewl rest of
x :< rest -> (x, noMatch >< rest)