Update base for latest Safe Haskell.
[packages/base.git] / GHC / Event / PSQ.hs
1 {-# LANGUAGE Trustworthy #-}
2 {-# LANGUAGE BangPatterns, NoImplicitPrelude #-}
3
4 -- Copyright (c) 2008, Ralf Hinze
5 -- All rights reserved.
6 --
7 -- Redistribution and use in source and binary forms, with or without
8 -- modification, are permitted provided that the following conditions
9 -- are met:
10 --
11 -- * Redistributions of source code must retain the above
12 -- copyright notice, this list of conditions and the following
13 -- disclaimer.
14 --
15 -- * Redistributions in binary form must reproduce the above
16 -- copyright notice, this list of conditions and the following
17 -- disclaimer in the documentation and/or other materials
18 -- provided with the distribution.
19 --
20 -- * The names of the contributors may not be used to endorse or
21 -- promote products derived from this software without specific
22 -- prior written permission.
23 --
24 -- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
25 -- "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
26 -- LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
27 -- FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
28 -- COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
29 -- INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
30 -- (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
31 -- SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
32 -- HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
33 -- STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
34 -- ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
35 -- OF THE POSSIBILITY OF SUCH DAMAGE.
36
37 -- | A /priority search queue/ (henceforth /queue/) efficiently
38 -- supports the operations of both a search tree and a priority queue.
39 -- An 'Elem'ent is a product of a key, a priority, and a
40 -- value. Elements can be inserted, deleted, modified and queried in
41 -- logarithmic time, and the element with the least priority can be
42 -- retrieved in constant time. A queue can be built from a list of
43 -- elements, sorted by keys, in linear time.
44 --
45 -- This implementation is due to Ralf Hinze with some modifications by
46 -- Scott Dillard and Johan Tibell.
47 --
48 -- * Hinze, R., /A Simple Implementation Technique for Priority Search
49 -- Queues/, ICFP 2001, pp. 110-121
50 --
51 -- <http://citeseer.ist.psu.edu/hinze01simple.html>
52 module GHC.Event.PSQ
53 (
54 -- * Binding Type
55 Elem(..)
56 , Key
57 , Prio
58
59 -- * Priority Search Queue Type
60 , PSQ
61
62 -- * Query
63 , size
64 , null
65 , lookup
66
67 -- * Construction
68 , empty
69 , singleton
70
71 -- * Insertion
72 , insert
73
74 -- * Delete/Update
75 , delete
76 , adjust
77
78 -- * Conversion
79 , toList
80 , toAscList
81 , toDescList
82 , fromList
83
84 -- * Min
85 , findMin
86 , deleteMin
87 , minView
88 , atMost
89 ) where
90
91 import Data.Maybe (Maybe(..))
92 import GHC.Base
93 import GHC.Num (Num(..))
94 import GHC.Show (Show(showsPrec))
95 import GHC.Event.Unique (Unique)
96
97 -- | @E k p@ binds the key @k@ with the priority @p@.
98 data Elem a = E
99 { key :: {-# UNPACK #-} !Key
100 , prio :: {-# UNPACK #-} !Prio
101 , value :: a
102 } deriving (Eq, Show)
103
104 ------------------------------------------------------------------------
105 -- | A mapping from keys @k@ to priorites @p@.
106
107 type Prio = Double
108 type Key = Unique
109
110 data PSQ a = Void
111 | Winner {-# UNPACK #-} !(Elem a)
112 !(LTree a)
113 {-# UNPACK #-} !Key -- max key
114 deriving (Eq, Show)
115
116 -- | /O(1)/ The number of elements in a queue.
117 size :: PSQ a -> Int
118 size Void = 0
119 size (Winner _ lt _) = 1 + size' lt
120
121 -- | /O(1)/ True if the queue is empty.
122 null :: PSQ a -> Bool
123 null Void = True
124 null (Winner _ _ _) = False
125
126 -- | /O(log n)/ The priority and value of a given key, or Nothing if
127 -- the key is not bound.
128 lookup :: Key -> PSQ a -> Maybe (Prio, a)
129 lookup k q = case tourView q of
130 Null -> Nothing
131 Single (E k' p v)
132 | k == k' -> Just (p, v)
133 | otherwise -> Nothing
134 tl `Play` tr
135 | k <= maxKey tl -> lookup k tl
136 | otherwise -> lookup k tr
137
138 ------------------------------------------------------------------------
139 -- Construction
140
141 empty :: PSQ a
142 empty = Void
143
144 -- | /O(1)/ Build a queue with one element.
145 singleton :: Key -> Prio -> a -> PSQ a
146 singleton k p v = Winner (E k p v) Start k
147
148 ------------------------------------------------------------------------
149 -- Insertion
150
151 -- | /O(log n)/ Insert a new key, priority and value in the queue. If
152 -- the key is already present in the queue, the associated priority
153 -- and value are replaced with the supplied priority and value.
154 insert :: Key -> Prio -> a -> PSQ a -> PSQ a
155 insert k p v q = case q of
156 Void -> singleton k p v
157 Winner (E k' p' v') Start _ -> case compare k k' of
158 LT -> singleton k p v `play` singleton k' p' v'
159 EQ -> singleton k p v
160 GT -> singleton k' p' v' `play` singleton k p v
161 Winner e (RLoser _ e' tl m tr) m'
162 | k <= m -> insert k p v (Winner e tl m) `play` (Winner e' tr m')
163 | otherwise -> (Winner e tl m) `play` insert k p v (Winner e' tr m')
164 Winner e (LLoser _ e' tl m tr) m'
165 | k <= m -> insert k p v (Winner e' tl m) `play` (Winner e tr m')
166 | otherwise -> (Winner e' tl m) `play` insert k p v (Winner e tr m')
167
168 ------------------------------------------------------------------------
169 -- Delete/Update
170
171 -- | /O(log n)/ Delete a key and its priority and value from the
172 -- queue. When the key is not a member of the queue, the original
173 -- queue is returned.
174 delete :: Key -> PSQ a -> PSQ a
175 delete k q = case q of
176 Void -> empty
177 Winner (E k' p v) Start _
178 | k == k' -> empty
179 | otherwise -> singleton k' p v
180 Winner e (RLoser _ e' tl m tr) m'
181 | k <= m -> delete k (Winner e tl m) `play` (Winner e' tr m')
182 | otherwise -> (Winner e tl m) `play` delete k (Winner e' tr m')
183 Winner e (LLoser _ e' tl m tr) m'
184 | k <= m -> delete k (Winner e' tl m) `play` (Winner e tr m')
185 | otherwise -> (Winner e' tl m) `play` delete k (Winner e tr m')
186
187 -- | /O(log n)/ Update a priority at a specific key with the result
188 -- of the provided function. When the key is not a member of the
189 -- queue, the original queue is returned.
190 adjust :: (Prio -> Prio) -> Key -> PSQ a -> PSQ a
191 adjust f k q0 = go q0
192 where
193 go q = case q of
194 Void -> empty
195 Winner (E k' p v) Start _
196 | k == k' -> singleton k' (f p) v
197 | otherwise -> singleton k' p v
198 Winner e (RLoser _ e' tl m tr) m'
199 | k <= m -> go (Winner e tl m) `unsafePlay` (Winner e' tr m')
200 | otherwise -> (Winner e tl m) `unsafePlay` go (Winner e' tr m')
201 Winner e (LLoser _ e' tl m tr) m'
202 | k <= m -> go (Winner e' tl m) `unsafePlay` (Winner e tr m')
203 | otherwise -> (Winner e' tl m) `unsafePlay` go (Winner e tr m')
204 {-# INLINE adjust #-}
205
206 ------------------------------------------------------------------------
207 -- Conversion
208
209 -- | /O(n*log n)/ Build a queue from a list of key/priority/value
210 -- tuples. If the list contains more than one priority and value for
211 -- the same key, the last priority and value for the key is retained.
212 fromList :: [Elem a] -> PSQ a
213 fromList = foldr (\(E k p v) q -> insert k p v q) empty
214
215 -- | /O(n)/ Convert to a list of key/priority/value tuples.
216 toList :: PSQ a -> [Elem a]
217 toList = toAscList
218
219 -- | /O(n)/ Convert to an ascending list.
220 toAscList :: PSQ a -> [Elem a]
221 toAscList q = seqToList (toAscLists q)
222
223 toAscLists :: PSQ a -> Sequ (Elem a)
224 toAscLists q = case tourView q of
225 Null -> emptySequ
226 Single e -> singleSequ e
227 tl `Play` tr -> toAscLists tl <> toAscLists tr
228
229 -- | /O(n)/ Convert to a descending list.
230 toDescList :: PSQ a -> [ Elem a ]
231 toDescList q = seqToList (toDescLists q)
232
233 toDescLists :: PSQ a -> Sequ (Elem a)
234 toDescLists q = case tourView q of
235 Null -> emptySequ
236 Single e -> singleSequ e
237 tl `Play` tr -> toDescLists tr <> toDescLists tl
238
239 ------------------------------------------------------------------------
240 -- Min
241
242 -- | /O(1)/ The element with the lowest priority.
243 findMin :: PSQ a -> Maybe (Elem a)
244 findMin Void = Nothing
245 findMin (Winner e _ _) = Just e
246
247 -- | /O(log n)/ Delete the element with the lowest priority. Returns
248 -- an empty queue if the queue is empty.
249 deleteMin :: PSQ a -> PSQ a
250 deleteMin Void = Void
251 deleteMin (Winner _ t m) = secondBest t m
252
253 -- | /O(log n)/ Retrieve the binding with the least priority, and the
254 -- rest of the queue stripped of that binding.
255 minView :: PSQ a -> Maybe (Elem a, PSQ a)
256 minView Void = Nothing
257 minView (Winner e t m) = Just (e, secondBest t m)
258
259 secondBest :: LTree a -> Key -> PSQ a
260 secondBest Start _ = Void
261 secondBest (LLoser _ e tl m tr) m' = Winner e tl m `play` secondBest tr m'
262 secondBest (RLoser _ e tl m tr) m' = secondBest tl m `play` Winner e tr m'
263
264 -- | /O(r*(log n - log r))/ Return a list of elements ordered by
265 -- key whose priorities are at most @pt@.
266 atMost :: Prio -> PSQ a -> ([Elem a], PSQ a)
267 atMost pt q = let (sequ, q') = atMosts pt q
268 in (seqToList sequ, q')
269
270 atMosts :: Prio -> PSQ a -> (Sequ (Elem a), PSQ a)
271 atMosts !pt q = case q of
272 (Winner e _ _)
273 | prio e > pt -> (emptySequ, q)
274 Void -> (emptySequ, Void)
275 Winner e Start _ -> (singleSequ e, Void)
276 Winner e (RLoser _ e' tl m tr) m' ->
277 let (sequ, q') = atMosts pt (Winner e tl m)
278 (sequ', q'') = atMosts pt (Winner e' tr m')
279 in (sequ <> sequ', q' `play` q'')
280 Winner e (LLoser _ e' tl m tr) m' ->
281 let (sequ, q') = atMosts pt (Winner e' tl m)
282 (sequ', q'') = atMosts pt (Winner e tr m')
283 in (sequ <> sequ', q' `play` q'')
284
285 ------------------------------------------------------------------------
286 -- Loser tree
287
288 type Size = Int
289
290 data LTree a = Start
291 | LLoser {-# UNPACK #-} !Size
292 {-# UNPACK #-} !(Elem a)
293 !(LTree a)
294 {-# UNPACK #-} !Key -- split key
295 !(LTree a)
296 | RLoser {-# UNPACK #-} !Size
297 {-# UNPACK #-} !(Elem a)
298 !(LTree a)
299 {-# UNPACK #-} !Key -- split key
300 !(LTree a)
301 deriving (Eq, Show)
302
303 size' :: LTree a -> Size
304 size' Start = 0
305 size' (LLoser s _ _ _ _) = s
306 size' (RLoser s _ _ _ _) = s
307
308 left, right :: LTree a -> LTree a
309
310 left Start = moduleError "left" "empty loser tree"
311 left (LLoser _ _ tl _ _ ) = tl
312 left (RLoser _ _ tl _ _ ) = tl
313
314 right Start = moduleError "right" "empty loser tree"
315 right (LLoser _ _ _ _ tr) = tr
316 right (RLoser _ _ _ _ tr) = tr
317
318 maxKey :: PSQ a -> Key
319 maxKey Void = moduleError "maxKey" "empty queue"
320 maxKey (Winner _ _ m) = m
321
322 lloser, rloser :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
323 lloser k p v tl m tr = LLoser (1 + size' tl + size' tr) (E k p v) tl m tr
324 rloser k p v tl m tr = RLoser (1 + size' tl + size' tr) (E k p v) tl m tr
325
326 ------------------------------------------------------------------------
327 -- Balancing
328
329 -- | Balance factor
330 omega :: Int
331 omega = 4
332
333 lbalance, rbalance :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
334
335 lbalance k p v l m r
336 | size' l + size' r < 2 = lloser k p v l m r
337 | size' r > omega * size' l = lbalanceLeft k p v l m r
338 | size' l > omega * size' r = lbalanceRight k p v l m r
339 | otherwise = lloser k p v l m r
340
341 rbalance k p v l m r
342 | size' l + size' r < 2 = rloser k p v l m r
343 | size' r > omega * size' l = rbalanceLeft k p v l m r
344 | size' l > omega * size' r = rbalanceRight k p v l m r
345 | otherwise = rloser k p v l m r
346
347 lbalanceLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
348 lbalanceLeft k p v l m r
349 | size' (left r) < size' (right r) = lsingleLeft k p v l m r
350 | otherwise = ldoubleLeft k p v l m r
351
352 lbalanceRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
353 lbalanceRight k p v l m r
354 | size' (left l) > size' (right l) = lsingleRight k p v l m r
355 | otherwise = ldoubleRight k p v l m r
356
357 rbalanceLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
358 rbalanceLeft k p v l m r
359 | size' (left r) < size' (right r) = rsingleLeft k p v l m r
360 | otherwise = rdoubleLeft k p v l m r
361
362 rbalanceRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
363 rbalanceRight k p v l m r
364 | size' (left l) > size' (right l) = rsingleRight k p v l m r
365 | otherwise = rdoubleRight k p v l m r
366
367 lsingleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
368 lsingleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3)
369 | p1 <= p2 = lloser k1 p1 v1 (rloser k2 p2 v2 t1 m1 t2) m2 t3
370 | otherwise = lloser k2 p2 v2 (lloser k1 p1 v1 t1 m1 t2) m2 t3
371 lsingleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
372 rloser k2 p2 v2 (lloser k1 p1 v1 t1 m1 t2) m2 t3
373 lsingleLeft _ _ _ _ _ _ = moduleError "lsingleLeft" "malformed tree"
374
375 rsingleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
376 rsingleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) =
377 rloser k1 p1 v1 (rloser k2 p2 v2 t1 m1 t2) m2 t3
378 rsingleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
379 rloser k2 p2 v2 (rloser k1 p1 v1 t1 m1 t2) m2 t3
380 rsingleLeft _ _ _ _ _ _ = moduleError "rsingleLeft" "malformed tree"
381
382 lsingleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
383 lsingleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
384 lloser k2 p2 v2 t1 m1 (lloser k1 p1 v1 t2 m2 t3)
385 lsingleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
386 lloser k1 p1 v1 t1 m1 (lloser k2 p2 v2 t2 m2 t3)
387 lsingleRight _ _ _ _ _ _ = moduleError "lsingleRight" "malformed tree"
388
389 rsingleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
390 rsingleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
391 lloser k2 p2 v2 t1 m1 (rloser k1 p1 v1 t2 m2 t3)
392 rsingleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3
393 | p1 <= p2 = rloser k1 p1 v1 t1 m1 (lloser k2 p2 v2 t2 m2 t3)
394 | otherwise = rloser k2 p2 v2 t1 m1 (rloser k1 p1 v1 t2 m2 t3)
395 rsingleRight _ _ _ _ _ _ = moduleError "rsingleRight" "malformed tree"
396
397 ldoubleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
398 ldoubleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) =
399 lsingleLeft k1 p1 v1 t1 m1 (lsingleRight k2 p2 v2 t2 m2 t3)
400 ldoubleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
401 lsingleLeft k1 p1 v1 t1 m1 (rsingleRight k2 p2 v2 t2 m2 t3)
402 ldoubleLeft _ _ _ _ _ _ = moduleError "ldoubleLeft" "malformed tree"
403
404 ldoubleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
405 ldoubleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
406 lsingleRight k1 p1 v1 (lsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
407 ldoubleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
408 lsingleRight k1 p1 v1 (rsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
409 ldoubleRight _ _ _ _ _ _ = moduleError "ldoubleRight" "malformed tree"
410
411 rdoubleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
412 rdoubleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) =
413 rsingleLeft k1 p1 v1 t1 m1 (lsingleRight k2 p2 v2 t2 m2 t3)
414 rdoubleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) =
415 rsingleLeft k1 p1 v1 t1 m1 (rsingleRight k2 p2 v2 t2 m2 t3)
416 rdoubleLeft _ _ _ _ _ _ = moduleError "rdoubleLeft" "malformed tree"
417
418 rdoubleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a
419 rdoubleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
420 rsingleRight k1 p1 v1 (lsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
421 rdoubleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 =
422 rsingleRight k1 p1 v1 (rsingleLeft k2 p2 v2 t1 m1 t2) m2 t3
423 rdoubleRight _ _ _ _ _ _ = moduleError "rdoubleRight" "malformed tree"
424
425 -- | Take two pennants and returns a new pennant that is the union of
426 -- the two with the precondition that the keys in the first tree are
427 -- strictly smaller than the keys in the second tree.
428 play :: PSQ a -> PSQ a -> PSQ a
429 Void `play` t' = t'
430 t `play` Void = t
431 Winner e@(E k p v) t m `play` Winner e'@(E k' p' v') t' m'
432 | p <= p' = Winner e (rbalance k' p' v' t m t') m'
433 | otherwise = Winner e' (lbalance k p v t m t') m'
434 {-# INLINE play #-}
435
436 -- | A version of 'play' that can be used if the shape of the tree has
437 -- not changed or if the tree is known to be balanced.
438 unsafePlay :: PSQ a -> PSQ a -> PSQ a
439 Void `unsafePlay` t' = t'
440 t `unsafePlay` Void = t
441 Winner e@(E k p v) t m `unsafePlay` Winner e'@(E k' p' v') t' m'
442 | p <= p' = Winner e (rloser k' p' v' t m t') m'
443 | otherwise = Winner e' (lloser k p v t m t') m'
444 {-# INLINE unsafePlay #-}
445
446 data TourView a = Null
447 | Single {-# UNPACK #-} !(Elem a)
448 | (PSQ a) `Play` (PSQ a)
449
450 tourView :: PSQ a -> TourView a
451 tourView Void = Null
452 tourView (Winner e Start _) = Single e
453 tourView (Winner e (RLoser _ e' tl m tr) m') =
454 Winner e tl m `Play` Winner e' tr m'
455 tourView (Winner e (LLoser _ e' tl m tr) m') =
456 Winner e' tl m `Play` Winner e tr m'
457
458 ------------------------------------------------------------------------
459 -- Utility functions
460
461 moduleError :: String -> String -> a
462 moduleError fun msg = error ("GHC.Event.PSQ." ++ fun ++ ':' : ' ' : msg)
463 {-# NOINLINE moduleError #-}
464
465 ------------------------------------------------------------------------
466 -- Hughes's efficient sequence type
467
468 newtype Sequ a = Sequ ([a] -> [a])
469
470 emptySequ :: Sequ a
471 emptySequ = Sequ (\as -> as)
472
473 singleSequ :: a -> Sequ a
474 singleSequ a = Sequ (\as -> a : as)
475
476 (<>) :: Sequ a -> Sequ a -> Sequ a
477 Sequ x1 <> Sequ x2 = Sequ (\as -> x1 (x2 as))
478 infixr 5 <>
479
480 seqToList :: Sequ a -> [a]
481 seqToList (Sequ x) = x []
482
483 instance Show a => Show (Sequ a) where
484 showsPrec d a = showsPrec d (seqToList a)
485