Skip to content
Fetching contributors…
Cannot retrieve contributors at this time
486 lines (406 sloc) 17.7 KB
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE BangPatterns, NoImplicitPrelude #-}
-- Copyright (c) 2008, Ralf Hinze
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions
-- are met:
--
-- * Redistributions of source code must retain the above
-- copyright notice, this list of conditions and the following
-- disclaimer.
--
-- * Redistributions in binary form must reproduce the above
-- copyright notice, this list of conditions and the following
-- disclaimer in the documentation and/or other materials
-- provided with the distribution.
--
-- * The names of the contributors may not be used to endorse or
-- promote products derived from this software without specific
-- prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
-- "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
-- LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
-- FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
-- COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
-- INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-- (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-- SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
-- HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
-- STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-- ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
-- OF THE POSSIBILITY OF SUCH DAMAGE.
-- | A /priority search queue/ (henceforth /queue/) efficiently
-- supports the operations of both a search tree and a priority queue.
-- An 'Elem'ent is a product of a key, a priority, and a
-- value. Elements can be inserted, deleted, modified and queried in
-- logarithmic time, and the element with the least priority can be
-- retrieved in constant time. A queue can be built from a list of
-- elements, sorted by keys, in linear time.
--
-- This implementation is due to Ralf Hinze with some modifications by
-- Scott Dillard and Johan Tibell.
--
-- * Hinze, R., /A Simple Implementation Technique for Priority Search
-- Queues/, ICFP 2001, pp. 110-121
--
-- <http://citeseer.ist.psu.edu/hinze01simple.html>
module GHC.Event.PSQ
(
-- * Binding Type
Elem(..)
, Key
, Prio
-- * Priority Search Queue Type
, PSQ
-- * Query
, size
, null
, lookup
-- * Construction
, empty
, singleton
-- * Insertion
, insert
-- * Delete/Update
, delete
, adjust
-- * Conversion
, toList
, toAscList
, toDescList
, fromList
-- * Min
, findMin
, deleteMin
, minView
, atMost
) where
import Data.Maybe (Maybe(..))
import GHC.Base
import GHC.Num (Num(..))
import GHC.Show (Show(showsPrec))
import GHC.Event.Unique (Unique)
-- | @E k p@ binds the key @k@ with the priority @p@.
data Elem a = E
{ key :: {-# UNPACK #-} !Key
, prio :: {-# UNPACK #-} !Prio
, value :: a
} deriving (Eq, Show)
------------------------------------------------------------------------
-- | A mapping from keys @k@ to priorites @p@.
type Prio = Double
type Key = Unique
data PSQ a = Void
| Winner {-# UNPACK #-} !(Elem a)
!(LTree a)
{-# UNPACK #-} !Key -- max key
deriving (Eq, Show)
-- | /O(1)/ The number of elements in a queue.
size :: PSQ a -> Int
size Void = 0
size (Winner _ lt _) = 1 + size' lt
-- | /O(1)/ True if the queue is empty.
null :: PSQ a -> Bool
null Void = True
null (Winner _ _ _) = False
-- | /O(log n)/ The priority and value of a given key, or Nothing if
-- the key is not bound.
lookup :: Key -> PSQ a -> Maybe (Prio, a)
lookup k q = case tourView q of
Null -> Nothing
Single (E k' p v)
| k == k' -> Just (p, v)
| otherwise -> Nothing
tl `Play` tr
| k <= maxKey tl -> lookup k tl
| otherwise -> lookup k tr
------------------------------------------------------------------------
-- Construction
empty :: PSQ a
empty = Void
-- | /O(1)/ Build a queue with one element.
singleton :: Key -> Prio -> a -> PSQ a
singleton k p v = Winner (E k p v) Start k
------------------------------------------------------------------------
-- Insertion
-- | /O(log n)/ Insert a new key, priority and value in the queue. If
-- the key is already present in the queue, the associated priority
-- and value are replaced with the supplied priority and value.
insert :: Key -> Prio -> a -> PSQ a -> PSQ a
insert k p v q = case q of
Void -> singleton k p v
Winner (E k' p' v') Start _ -> case compare k k' of
LT -> singleton k p v `play` singleton k' p' v'
EQ -> singleton k p v
GT -> singleton k' p' v' `play` singleton k p v
Winner e (RLoser _ e' tl m tr) m'
| k <= m -> insert k p v (Winner e tl m) `play` (Winner e' tr m')
| otherwise -> (Winner e tl m) `play` insert k p v (Winner e' tr m')
Winner e (LLoser _ e' tl m tr) m'
| k <= m -> insert k p v (Winner e' tl m) `play` (Winner e tr m')
| otherwise -> (Winner e' tl m) `play` insert k p v (Winner e tr m')
------------------------------------------------------------------------
-- Delete/Update
-- | /O(log n)/ Delete a key and its priority and value from the
-- queue. When the key is not a member of the queue, the original
-- queue is returned.
delete :: Key -> PSQ a -> PSQ a
delete k q = case q of
Void -> empty
Winner (E k' p v) Start _
| k == k' -> empty
| otherwise -> singleton k' p v
Winner e (RLoser _ e' tl m tr) m'
| k <= m -> delete k (Winner e tl m) `play` (Winner e' tr m')
| otherwise -> (Winner e tl m) `play` delete k (Winner e' tr m')
Winner e (LLoser _ e' tl m tr) m'
| k <= m -> delete k (Winner e' tl m) `play` (Winner e tr m')
| otherwise -> (Winner e' tl m) `play` delete k (Winner e tr m')
-- | /O(log n)/ Update a priority at a specific key with the result
-- of the provided function. When the key is not a member of the
-- queue, the original queue is returned.
adjust :: (Prio -> Prio) -> Key -> PSQ a -> PSQ a
adjust f k q0 = go q0
where
go q = case q of
Void -> empty
Winner (E k' p v) Start _
| k == k' -> singleton k' (f p) v
| otherwise -> singleton k' p v
Winner e (RLoser _ e' tl m tr) m'
| k <= m -> go (Winner e tl m) `unsafePlay` (Winner e' tr m')
| otherwise -> (Winner e tl m) `unsafePlay` go (Winner e' tr m')
Winner e (LLoser _ e' tl m tr) m'
| k <= m -> go (Winner e' tl m) `unsafePlay` (Winner e tr m')
| otherwise -> (Winner e' tl m) `unsafePlay` go (Winner e tr m')
{-# INLINE adjust #-}
------------------------------------------------------------------------
-- Conversion
-- | /O(n*log n)/ Build a queue from a list of key/priority/value
-- tuples. If the list contains more than one priority and value for
-- the same key, the last priority and value for the key is retained.
fromList :: [Elem a] -> PSQ a
fromList = foldr (\(E k p v) q -> insert k p v q) empty
-- | /O(n)/ Convert to a list of key/priority/value tuples.
toList :: PSQ a -> [Elem a]
toList = toAscList
-- | /O(n)/ Convert to an ascending list.
toAscList :: PSQ a -> [Elem a]
toAscList q = seqToList (toAscLists q)
toAscLists :: PSQ a -> Sequ (Elem a)
toAscLists q = case tourView q of
Null -> emptySequ
Single e -> singleSequ e
tl `Play` tr -> toAscLists tl <> toAscLists tr
-- | /O(n)/ Convert to a descending list.
toDescList :: PSQ a -> [ Elem a ]
toDescList q = seqToList (toDescLists q)
toDescLists :: PSQ a -> Sequ (Elem a)
toDescLists q = case tourView q of
Null -> emptySequ
Single e -> singleSequ e
tl `Play` tr -> toDescLists tr <> toDescLists tl
------------------------------------------------------------------------
-- Min
-- | /O(1)/ The element with the lowest priority.
findMin :: PSQ a -> Maybe (Elem a)
findMin Void = Nothing
findMin (Winner e _ _) = Just e
-- | /O(log n)/ Delete the element with the lowest priority. Returns
-- an empty queue if the queue is empty.
deleteMin :: PSQ a -> PSQ a
deleteMin Void = Void
deleteMin (Winner _ t m) = secondBest t m
-- | /O(log n)/ Retrieve the binding with the least priority, and the
-- rest of the queue stripped of that binding.
minView :: PSQ a -> Maybe (Elem a, PSQ a)
minView Void = Nothing
minView (Winner e t m) = Just (e, secondBest t m)
secondBest :: LTree a -> Key -> PSQ a
secondBest Start _ = Void
secondBest (LLoser _ e tl m tr) m' = Winner e tl m `play` secondBest tr m'
secondBest (RLoser _ e tl m tr) m' = secondBest tl m `play` Winner e tr m'
-- | /O(r*(log n - log r))/ Return a list of elements ordered by
-- key whose priorities are at most @pt@.
atMost :: Prio -> PSQ a -> ([Elem a], PSQ a)
atMost pt q = let (sequ, q') = atMosts pt q
in (seqToList sequ, q')
atMosts :: Prio -> PSQ a -> (Sequ (Elem a), PSQ a)
atMosts !pt q = case q of
(Winner e _ _)
| prio e > pt -> (emptySequ, q)
Void -> (emptySequ, Void)
Winner e Start _ -> (singleSequ e, Void)
Winner e (RLoser _ e' tl m tr) m' ->
let (sequ, q') = atMosts pt (Winner e tl m)
(sequ', q'') = atMosts pt (Winner e' tr m')
in (sequ <> sequ', q' `play` q'')
Winner e (LLoser _ e' tl m tr) m' ->
let (sequ, q') = atMosts pt (Winner e' tl m)
(sequ', q'') = atMosts pt (Winner e tr m')
in (sequ <> sequ', q' `play` q'')
------------------------------------------------------------------------
-- Loser tree
type Size = Int
data LTree a = Start
| LLoser {-# UNPACK #-} !Size
{-# UNPACK #-} !(Elem a)
!(LTree a)
{-# UNPACK #-} !Key -- split key
!(LTree a)
| RLoser {-# UNPACK #-} !Size
{-# UNPACK #-} !(Elem a)
!(LTree a)
{-# UNPACK #-} !Key -- split key
!(LTree a)
deriving (Eq, Show)
size' :: LTree a -> Size
size' Start = 0
size' (LLoser s _ _ _ _) = s
size' (RLoser s _ _ _ _) = s
left, right :: LTree a -> LTree a
left Start = moduleError "left" "empty loser tree"
left (LLoser _ _ tl _ _ ) = tl
left (RLoser _ _ tl _ _ ) = tl
right Start = moduleError "right" "empty loser tree"
right (LLoser _ _ _ _ tr) = tr
right (RLoser _ _ _ _ tr) = tr
maxKey :: PSQ a -> Key
maxKey Void = moduleError "maxKey" "empty queue"
maxKey (Winner _ _ m) = m
lloser, rloser :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
lloser k p v tl m tr = LLoser (1 + size' tl + size' tr) (E k p v) tl m tr
rloser k p v tl m tr = RLoser (1 + size' tl + size' tr) (E k p v) tl m tr
------------------------------------------------------------------------
-- Balancing
-- | Balance factor
omega :: Int
omega = 4
lbalance, rbalance :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
lbalance k p v l m r
| size' l + size' r < 2 = lloser k p v l m r
| size' r > omega * size' l = lbalanceLeft k p v l m r
| size' l > omega * size' r = lbalanceRight k p v l m r
| otherwise = lloser k p v l m r
rbalance k p v l m r
| size' l + size' r < 2 = rloser k p v l m r
| size' r > omega * size' l = rbalanceLeft k p v l m r
| size' l > omega * size' r = rbalanceRight k p v l m r
| otherwise = rloser k p v l m r
lbalanceLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
lbalanceLeft k p v l m r
| size' (left r) < size' (right r) = lsingleLeft k p v l m r
| otherwise = ldoubleLeft k p v l m r
lbalanceRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
lbalanceRight k p v l m r
| size' (left l) > size' (right l) = lsingleRight k p v l m r
| otherwise = ldoubleRight k p v l m r
rbalanceLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
rbalanceLeft k p v l m r
| size' (left r) < size' (right r) = rsingleLeft k p v l m r
| otherwise = rdoubleLeft k p v l m r
rbalanceRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
rbalanceRight k p v l m r
| size' (left l) > size' (right l) = rsingleRight k p v l m r
| otherwise = rdoubleRight k p v l m r
lsingleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
lsingleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3)
| p1 <= p2 = lloser k1 p1 v1 (rloser k2 p2 v2 t1 m1 t2) m2 t3
| otherwise = lloser k2 p2 v2 (lloser k1 p1 v1 t1 m1 t2) m2 t3
lsingleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
rloser k2 p2 v2 (lloser k1 p1 v1 t1 m1 t2) m2 t3
lsingleLeft _ _ _ _ _ _ = moduleError "lsingleLeft" "malformed tree"
rsingleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
rsingleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) =
rloser k1 p1 v1 (rloser k2 p2 v2 t1 m1 t2) m2 t3
rsingleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
rloser k2 p2 v2 (rloser k1 p1 v1 t1 m1 t2) m2 t3
rsingleLeft _ _ _ _ _ _ = moduleError "rsingleLeft" "malformed tree"
lsingleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
lsingleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
lloser k2 p2 v2 t1 m1 (lloser k1 p1 v1 t2 m2 t3)
lsingleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
lloser k1 p1 v1 t1 m1 (lloser k2 p2 v2 t2 m2 t3)
lsingleRight _ _ _ _ _ _ = moduleError "lsingleRight" "malformed tree"
rsingleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
rsingleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
lloser k2 p2 v2 t1 m1 (rloser k1 p1 v1 t2 m2 t3)
rsingleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3
| p1 <= p2 = rloser k1 p1 v1 t1 m1 (lloser k2 p2 v2 t2 m2 t3)
| otherwise = rloser k2 p2 v2 t1 m1 (rloser k1 p1 v1 t2 m2 t3)
rsingleRight _ _ _ _ _ _ = moduleError "rsingleRight" "malformed tree"
ldoubleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
ldoubleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) =
lsingleLeft k1 p1 v1 t1 m1 (lsingleRight k2 p2 v2 t2 m2 t3)
ldoubleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
lsingleLeft k1 p1 v1 t1 m1 (rsingleRight k2 p2 v2 t2 m2 t3)
ldoubleLeft _ _ _ _ _ _ = moduleError "ldoubleLeft" "malformed tree"
ldoubleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
ldoubleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
lsingleRight k1 p1 v1 (lsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
ldoubleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
lsingleRight k1 p1 v1 (rsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
ldoubleRight _ _ _ _ _ _ = moduleError "ldoubleRight" "malformed tree"
rdoubleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
rdoubleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) =
rsingleLeft k1 p1 v1 t1 m1 (lsingleRight k2 p2 v2 t2 m2 t3)
rdoubleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
rsingleLeft k1 p1 v1 t1 m1 (rsingleRight k2 p2 v2 t2 m2 t3)
rdoubleLeft _ _ _ _ _ _ = moduleError "rdoubleLeft" "malformed tree"
rdoubleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
rdoubleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
rsingleRight k1 p1 v1 (lsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
rdoubleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
rsingleRight k1 p1 v1 (rsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
rdoubleRight _ _ _ _ _ _ = moduleError "rdoubleRight" "malformed tree"
-- | Take two pennants and returns a new pennant that is the union of
-- the two with the precondition that the keys in the first tree are
-- strictly smaller than the keys in the second tree.
play :: PSQ a -> PSQ a -> PSQ a
Void `play` t' = t'
t `play` Void = t
Winner e@(E k p v) t m `play` Winner e'@(E k' p' v') t' m'
| p <= p' = Winner e (rbalance k' p' v' t m t') m'
| otherwise = Winner e' (lbalance k p v t m t') m'
{-# INLINE play #-}
-- | A version of 'play' that can be used if the shape of the tree has
-- not changed or if the tree is known to be balanced.
unsafePlay :: PSQ a -> PSQ a -> PSQ a
Void `unsafePlay` t' = t'
t `unsafePlay` Void = t
Winner e@(E k p v) t m `unsafePlay` Winner e'@(E k' p' v') t' m'
| p <= p' = Winner e (rloser k' p' v' t m t') m'
| otherwise = Winner e' (lloser k p v t m t') m'
{-# INLINE unsafePlay #-}
data TourView a = Null
| Single {-# UNPACK #-} !(Elem a)
| (PSQ a) `Play` (PSQ a)
tourView :: PSQ a -> TourView a
tourView Void = Null
tourView (Winner e Start _) = Single e
tourView (Winner e (RLoser _ e' tl m tr) m') =
Winner e tl m `Play` Winner e' tr m'
tourView (Winner e (LLoser _ e' tl m tr) m') =
Winner e' tl m `Play` Winner e tr m'
------------------------------------------------------------------------
-- Utility functions
moduleError :: String -> String -> a
moduleError fun msg = error ("GHC.Event.PSQ." ++ fun ++ ':' : ' ' : msg)
{-# NOINLINE moduleError #-}
------------------------------------------------------------------------
-- Hughes's efficient sequence type
newtype Sequ a = Sequ ([a] -> [a])
emptySequ :: Sequ a
emptySequ = Sequ (\as -> as)
singleSequ :: a -> Sequ a
singleSequ a = Sequ (\as -> a : as)
(<>) :: Sequ a -> Sequ a -> Sequ a
Sequ x1 <> Sequ x2 = Sequ (\as -> x1 (x2 as))
infixr 5 <>
seqToList :: Sequ a -> [a]
seqToList (Sequ x) = x []
instance Show a => Show (Sequ a) where
showsPrec d a = showsPrec d (seqToList a)
Something went wrong with that request. Please try again.