PriorityQueue.hs

{-
 -      ``Data/PriorityQueue''
 -      (c) 2008 Cook, J. MR  SSD, Inc.
 -      
 -      the PriorityQueue kicks ass, if I do say so myself ;-)
 -      the |DefaultStateRef| class makes the choice of StateRef
 -      decidable, and the laxity of the StateRef classes' fundeps makes
 -      queues constructible in monads other than where they are intended
 -      to be used; eg:
 -      
 -         q <- newPriorityQueue show :: IO (PriorityQueue STM Integer)
 -      
 -      after which the whole interface to the queue is:
 -         enqueue (x :: Integer) q :: STM ()
 -         dequeue q :: STM Integer
 -      
 -      If the queue is being constructed in the same scope it is used,
 -      the full type of |newPriorityQueue f| can be inferred as well,
 -      as long as |f|'s target type is monomorphic.
 -      
 -}
{-# LANGUAGE
        ExistentialQuantification,
        MultiParamTypeClasses,
        FlexibleContexts,
        FlexibleInstances,
        CPP
  #-}

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

-- |The "pure" type at the chewy center.
data PQ a = forall p. Ord p =>
        PQ { priorityFunc       :: a -> p
           , queue              :: M.Map p [a]
           }

-- |A new empty 'PQ'
emptyPQ :: Ord p => (a -> p) -> PQ a
emptyPQ f = PQ f M.empty

-- |A priority queue usable in the monad 'm' with values of type 'a'
data PriorityQueue m a = 
    forall sr. ( ModifyRef sr m (PQ a) 
               ) => PriorityQueue sr

-- |Build a priority queue from a modifiable reference containing
--  a 'PQ'
mkPriorityQueue :: ModifyRef sr m (PQ a) => sr -> PriorityQueue m a
mkPriorityQueue = PriorityQueue

-- |Build a priority queue using an instance of the default modifiable 
--  reference for the requested monad and value type
mkDefaultPriorityQueue :: 
        ( DefaultStateRef sr m (PQ a)
        , ModifyRef sr m (PQ a)
        ) => sr -> PriorityQueue m a
mkDefaultPriorityQueue = PriorityQueue

-- |Construct a new priority queue using the specified indexing function
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)

-- |Construct a new priority queue using a comparator function.  It is 
--  the user's responsibility to ensure that this function provides a
--  sensible order.
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')      -> 
                                        -- 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) = atomicModifyRef 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)

-- 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 = 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)

-- |local version of minViewWithKey, because some versions of Data.Map
--  don't have it.
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