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 Control.Arrow ((&&&), (***))
import Data.StateRef
import Data.Ord.ReOrd
import qualified Data.Map as M
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 (Seq 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 :: Ref m (PQ a) -> PriorityQueue m a
mkDefaultPriorityQueue = PriorityQueue

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

-- |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 :: 
        ( 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)
    
    -- 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 (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 $ \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)

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

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