diff --git a/priority-queue.cabal b/priority-queue.cabal --- a/priority-queue.cabal +++ b/priority-queue.cabal @@ -1,5 +1,5 @@ name: priority-queue -version: 0.2 +version: 0.2.1 stability: provisional license: BSD3 license-file: LICENSE @@ -30,4 +30,4 @@ cpp-options: -DNoMinViewWithKey build-depends: reord >= 0.0.0.2, stateref >= 0.3 && < 0.4, - queue >= 0.1.1.3 + queue >= 0.1.2 && < 0.2 diff --git a/src/Data/PriorityQueue.hs b/src/Data/PriorityQueue.hs --- a/src/Data/PriorityQueue.hs +++ b/src/Data/PriorityQueue.hs @@ -47,15 +47,18 @@ import Data.Queue.Classes +import Control.Arrow ((&&&), (***)) import Data.StateRef import Data.Ord.ReOrd import qualified Data.Map as M -import Data.List +import Data.Sequence as Seq +import Data.List as List +import Data.Foldable as Foldable -- |The "pure" type at the chewy center. data PQ a = forall p. Ord p => PQ { priorityFunc :: a -> p - , queue :: M.Map p [a] + , queue :: M.Map p (Seq a) } -- |A new empty 'PQ' @@ -100,55 +103,62 @@ newPriorityQueueBy cmp = newPriorityQueue (ReOrd cmp) instance Monad m => Enqueue (PriorityQueue m a) m a where - enqueue (PriorityQueue pqRef) x = modifyReference pqRef ins - where ins (PQ f pq) = PQ f (M.insertWith (flip (++)) (f x) [x] pq) + enqueue (PriorityQueue pqRef) x = modifyReference pqRef $ \(PQ f pq) -> + PQ f (M.insertWith (flip (><)) (f x) (singleton x) pq) + + -- the presumption here is that this is normally called for a bunch of + -- elements of the same priority, so we prepare the input list by + -- grouping elements by priority. In cases where the batch does have + -- large blocks of elements with the same priority, this will greatly + -- reduce the amount of work done by 'M.fromListWith'. TODO: Test + -- whether (and when) this is worth the extra initial traversal. Also + -- check to make sure as much list fusion as I expect is actually + -- happening. + 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 q@(PriorityQueue pqRef) = atomicModifyReference pqRef dq - where - dq orig@(PQ f pq) = case minViewWithKey pq of - Nothing -> (orig, Nothing) - Just ((k,[]), pq') -> - -- this should never happen - 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) = atomicModifyReference pqRef dq - where - dq orig@(PQ f pq) = case M.minView pq of - Nothing -> (orig, []) - Just ([], pq') -> - -- this should never happen - dq (PQ f pq') - Just (xs, pq') -> (PQ f pq', xs) + 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) + -- quick hack; there's probably a more efficient (and/or less ugly) way to do this instance Monad m => DequeueWhere (PriorityQueue m a) m a where - dequeueWhere (PriorityQueue pqRef) p = atomicModifyReference 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) + 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 <- vs] + 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.size pq) + queueSize (PriorityQueue pqRef) = do + PQ f pq <- readReference pqRef + return (M.fold (\xs t -> Seq.length xs + t) 0 pq) -- |local version of minViewWithKey, because some versions of Data.Map -- don't have it. @@ -161,3 +171,14 @@ #else minViewWithKey = M.minViewWithKey #endif + +-- |'on' combinator (Data.Function doesn't always have it) +on :: (b -> b -> c) -> (a -> b) -> (a -> a -> c) +op `on` f = \x y -> f x `op` f y + +-- |given a Seq known to contain at least one item matching the predicate, +-- return the (first) matching item and the seq sans that element +extractFirstWhere :: (a -> Bool) -> Seq a -> (a, Seq a) +extractFirstWhere p xs = case Seq.breakl p xs of + (noMatch, rest) -> case viewl rest of + x :< rest -> (x, noMatch >< rest)