priority-queue 0.2 → 0.2.1
raw patch · 2 files changed
+65/−44 lines, 2 filesdep ~queue
Dependency ranges changed: queue
Files
- priority-queue.cabal +2/−2
- src/Data/PriorityQueue.hs +63/−42
priority-queue.cabal view
@@ -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
src/Data/PriorityQueue.hs view
@@ -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)