Implement more methods from Foldable class.
[packages/containers.git] / Data / IntMap.hs
1 {-# LANGUAGE CPP, NoBangPatterns, MagicHash, ScopedTypeVariables #-}
2 -----------------------------------------------------------------------------
3 -- |
4 -- Module : Data.IntMap
5 -- Copyright : (c) Daan Leijen 2002
6 -- (c) Andriy Palamarchuk 2008
7 -- License : BSD-style
8 -- Maintainer : libraries@haskell.org
9 -- Stability : provisional
10 -- Portability : portable
11 --
12 -- An efficient implementation of maps from integer keys to values.
13 --
14 -- Since many function names (but not the type name) clash with
15 -- "Prelude" names, this module is usually imported @qualified@, e.g.
16 --
17 -- > import Data.IntMap (IntMap)
18 -- > import qualified Data.IntMap as IntMap
19 --
20 -- The implementation is based on /big-endian patricia trees/. This data
21 -- structure performs especially well on binary operations like 'union'
22 -- and 'intersection'. However, my benchmarks show that it is also
23 -- (much) faster on insertions and deletions when compared to a generic
24 -- size-balanced map implementation (see "Data.Map").
25 --
26 -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
27 -- Workshop on ML, September 1998, pages 77-86,
28 -- <http://citeseer.ist.psu.edu/okasaki98fast.html>
29 --
30 -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
31 -- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
32 -- October 1968, pages 514-534.
33 --
34 -- Operation comments contain the operation time complexity in
35 -- the Big-O notation <http://en.wikipedia.org/wiki/Big_O_notation>.
36 -- Many operations have a worst-case complexity of /O(min(n,W))/.
37 -- This means that the operation can become linear in the number of
38 -- elements with a maximum of /W/ -- the number of bits in an 'Int'
39 -- (32 or 64).
40 -----------------------------------------------------------------------------
41
42 -- It is essential that the bit fiddling functions like mask, zero, branchMask
43 -- etc are inlined. If they do not, the memory allocation skyrockets. The GHC
44 -- usually gets it right, but it is disastrous if it does not. Therefore we
45 -- explicitly mark these functions INLINE.
46
47 module Data.IntMap (
48 -- * Map type
49 #if !defined(TESTING)
50 IntMap, Key -- instance Eq,Show
51 #else
52 IntMap(..), Key -- instance Eq,Show
53 #endif
54
55 -- * Operators
56 , (!), (\\)
57
58 -- * Query
59 , null
60 , size
61 , member
62 , notMember
63 , lookup
64 , findWithDefault
65
66 -- * Construction
67 , empty
68 , singleton
69
70 -- ** Insertion
71 , insert
72 , insertWith
73 , insertWith'
74 , insertWithKey
75 , insertWithKey'
76 , insertLookupWithKey
77
78 -- ** Delete\/Update
79 , delete
80 , adjust
81 , adjustWithKey
82 , update
83 , updateWithKey
84 , updateLookupWithKey
85 , alter
86
87 -- * Combine
88
89 -- ** Union
90 , union
91 , unionWith
92 , unionWithKey
93 , unions
94 , unionsWith
95
96 -- ** Difference
97 , difference
98 , differenceWith
99 , differenceWithKey
100
101 -- ** Intersection
102 , intersection
103 , intersectionWith
104 , intersectionWithKey
105
106 -- * Traversal
107 -- ** Map
108 , map
109 , mapWithKey
110 , mapAccum
111 , mapAccumWithKey
112 , mapAccumRWithKey
113
114 -- * Folds
115 , foldr
116 , foldl
117 , foldrWithKey
118 , foldlWithKey
119 -- ** Strict folds
120 , foldr'
121 , foldl'
122 , foldrWithKey'
123 , foldlWithKey'
124 -- ** Legacy folds
125 , fold
126 , foldWithKey
127
128 -- * Conversion
129 , elems
130 , keys
131 , keysSet
132 , assocs
133
134 -- ** Lists
135 , toList
136 , fromList
137 , fromListWith
138 , fromListWithKey
139
140 -- ** Ordered lists
141 , toAscList
142 , fromAscList
143 , fromAscListWith
144 , fromAscListWithKey
145 , fromDistinctAscList
146
147 -- * Filter
148 , filter
149 , filterWithKey
150 , partition
151 , partitionWithKey
152
153 , mapMaybe
154 , mapMaybeWithKey
155 , mapEither
156 , mapEitherWithKey
157
158 , split
159 , splitLookup
160
161 -- * Submap
162 , isSubmapOf, isSubmapOfBy
163 , isProperSubmapOf, isProperSubmapOfBy
164
165 -- * Min\/Max
166 , findMin
167 , findMax
168 , deleteMin
169 , deleteMax
170 , deleteFindMin
171 , deleteFindMax
172 , updateMin
173 , updateMax
174 , updateMinWithKey
175 , updateMaxWithKey
176 , minView
177 , maxView
178 , minViewWithKey
179 , maxViewWithKey
180
181 -- * Debugging
182 , showTree
183 , showTreeWith
184 ) where
185
186 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
187 import Data.Bits
188 import qualified Data.IntSet as IntSet
189 import Data.Monoid (Monoid(..))
190 import Data.Maybe (fromMaybe)
191 import Data.Typeable
192 import qualified Data.Foldable as Foldable
193 import Data.Traversable (Traversable(traverse))
194 import Control.Applicative (Applicative(pure,(<*>)),(<$>))
195 import Control.Monad ( liftM )
196 {-
197 -- just for testing
198 import qualified Prelude
199 import Test.QuickCheck
200 import List (nub,sort)
201 import qualified List
202 -}
203
204 #if __GLASGOW_HASKELL__
205 import Text.Read
206 import Data.Data (Data(..), mkNoRepType)
207 #endif
208
209 #if __GLASGOW_HASKELL__ >= 503
210 import GHC.Exts ( Word(..), Int(..), shiftRL# )
211 #elif __GLASGOW_HASKELL__
212 import Word
213 import GlaExts ( Word(..), Int(..), shiftRL# )
214 #else
215 import Data.Word
216 #endif
217
218 -- Use macros to define strictness of functions.
219 -- STRICT_x_OF_y denotes an y-ary function strict in the x-th parameter.
220 -- We do not use BangPatterns, because they are not in any standard and we
221 -- want the compilers to be compiled by as many compilers as possible.
222 #define STRICT_1_OF_2(fn) fn arg _ | arg `seq` False = undefined
223
224 infixl 9 \\{-This comment teaches CPP correct behaviour -}
225
226 -- A "Nat" is a natural machine word (an unsigned Int)
227 type Nat = Word
228
229 natFromInt :: Key -> Nat
230 natFromInt = fromIntegral
231 {-# INLINE natFromInt #-}
232
233 intFromNat :: Nat -> Key
234 intFromNat = fromIntegral
235 {-# INLINE intFromNat #-}
236
237 shiftRL :: Nat -> Key -> Nat
238 #if __GLASGOW_HASKELL__
239 {--------------------------------------------------------------------
240 GHC: use unboxing to get @shiftRL@ inlined.
241 --------------------------------------------------------------------}
242 shiftRL (W# x) (I# i)
243 = W# (shiftRL# x i)
244 #else
245 shiftRL x i = shiftR x i
246 {-# INLINE shiftRL #-}
247 #endif
248
249 {--------------------------------------------------------------------
250 Operators
251 --------------------------------------------------------------------}
252
253 -- | /O(min(n,W))/. Find the value at a key.
254 -- Calls 'error' when the element can not be found.
255 --
256 -- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
257 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
258
259 (!) :: IntMap a -> Key -> a
260 m ! k = find k m
261
262 -- | Same as 'difference'.
263 (\\) :: IntMap a -> IntMap b -> IntMap a
264 m1 \\ m2 = difference m1 m2
265
266 {--------------------------------------------------------------------
267 Types
268 --------------------------------------------------------------------}
269
270 -- The order of constructors of IntMap matters when considering performance.
271 -- Currently in GHC 7.0, when type has 3 constructors, they are matched from
272 -- the first to the last -- the best performance is achieved when the
273 -- constructors are ordered by frequency.
274 -- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil
275 -- improves the containers_benchmark by 9.5% on x86 and by 8% on x86_64.
276
277 -- | A map of integers to values @a@.
278 data IntMap a = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
279 | Tip {-# UNPACK #-} !Key a
280 | Nil
281
282 type Prefix = Int
283 type Mask = Int
284 type Key = Int
285
286 instance Monoid (IntMap a) where
287 mempty = empty
288 mappend = union
289 mconcat = unions
290
291 instance Foldable.Foldable IntMap where
292 fold Nil = mempty
293 fold (Tip _ v) = v
294 fold (Bin _ _ l r) = Foldable.fold l `mappend` Foldable.fold r
295 foldr = foldr
296 foldl = foldl
297 foldMap _ Nil = mempty
298 foldMap f (Tip _k v) = f v
299 foldMap f (Bin _ _ l r) = Foldable.foldMap f l `mappend` Foldable.foldMap f r
300
301 instance Traversable IntMap where
302 traverse _ Nil = pure Nil
303 traverse f (Tip k v) = Tip k <$> f v
304 traverse f (Bin p m l r) = Bin p m <$> traverse f l <*> traverse f r
305
306 #if __GLASGOW_HASKELL__
307
308 {--------------------------------------------------------------------
309 A Data instance
310 --------------------------------------------------------------------}
311
312 -- This instance preserves data abstraction at the cost of inefficiency.
313 -- We omit reflection services for the sake of data abstraction.
314
315 instance Data a => Data (IntMap a) where
316 gfoldl f z im = z fromList `f` (toList im)
317 toConstr _ = error "toConstr"
318 gunfold _ _ = error "gunfold"
319 dataTypeOf _ = mkNoRepType "Data.IntMap.IntMap"
320 dataCast1 f = gcast1 f
321
322 #endif
323
324 {--------------------------------------------------------------------
325 Query
326 --------------------------------------------------------------------}
327 -- | /O(1)/. Is the map empty?
328 --
329 -- > Data.IntMap.null (empty) == True
330 -- > Data.IntMap.null (singleton 1 'a') == False
331
332 null :: IntMap a -> Bool
333 null Nil = True
334 null _ = False
335
336 -- | /O(n)/. Number of elements in the map.
337 --
338 -- > size empty == 0
339 -- > size (singleton 1 'a') == 1
340 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
341 size :: IntMap a -> Int
342 size t
343 = case t of
344 Bin _ _ l r -> size l + size r
345 Tip _ _ -> 1
346 Nil -> 0
347
348 -- | /O(min(n,W))/. Is the key a member of the map?
349 --
350 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
351 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
352
353 member :: Key -> IntMap a -> Bool
354 member k m
355 = case lookup k m of
356 Nothing -> False
357 Just _ -> True
358
359 -- | /O(log n)/. Is the key not a member of the map?
360 --
361 -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
362 -- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
363
364 notMember :: Key -> IntMap a -> Bool
365 notMember k m = not $ member k m
366
367 -- The 'go' function in the lookup causes 10% speedup, but also an increased
368 -- memory allocation. It does not cause speedup with other methods like insert
369 -- and delete, so it is present only in lookup.
370
371 -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'.
372 lookup :: Key -> IntMap a -> Maybe a
373 lookup k = k `seq` go
374 where
375 go (Bin _ m l r)
376 | zero k m = go l
377 | otherwise = go r
378 go (Tip kx x)
379 | k == kx = Just x
380 | otherwise = Nothing
381 go Nil = Nothing
382
383
384 find :: Key -> IntMap a -> a
385 find k m
386 = case lookup k m of
387 Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
388 Just x -> x
389
390 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
391 -- returns the value at key @k@ or returns @def@ when the key is not an
392 -- element of the map.
393 --
394 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
395 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
396
397 findWithDefault :: a -> Key -> IntMap a -> a
398 findWithDefault def k m
399 = case lookup k m of
400 Nothing -> def
401 Just x -> x
402
403 {--------------------------------------------------------------------
404 Construction
405 --------------------------------------------------------------------}
406 -- | /O(1)/. The empty map.
407 --
408 -- > empty == fromList []
409 -- > size empty == 0
410
411 empty :: IntMap a
412 empty
413 = Nil
414
415 -- | /O(1)/. A map of one element.
416 --
417 -- > singleton 1 'a' == fromList [(1, 'a')]
418 -- > size (singleton 1 'a') == 1
419
420 singleton :: Key -> a -> IntMap a
421 singleton k x
422 = Tip k x
423
424 {--------------------------------------------------------------------
425 Insert
426 --------------------------------------------------------------------}
427 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
428 -- If the key is already present in the map, the associated value is
429 -- replaced with the supplied value, i.e. 'insert' is equivalent to
430 -- @'insertWith' 'const'@.
431 --
432 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
433 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
434 -- > insert 5 'x' empty == singleton 5 'x'
435
436 insert :: Key -> a -> IntMap a -> IntMap a
437 insert k x t = k `seq`
438 case t of
439 Bin p m l r
440 | nomatch k p m -> join k (Tip k x) p t
441 | zero k m -> Bin p m (insert k x l) r
442 | otherwise -> Bin p m l (insert k x r)
443 Tip ky _
444 | k==ky -> Tip k x
445 | otherwise -> join k (Tip k x) ky t
446 Nil -> Tip k x
447
448 -- right-biased insertion, used by 'union'
449 -- | /O(min(n,W))/. Insert with a combining function.
450 -- @'insertWith' f key value mp@
451 -- will insert the pair (key, value) into @mp@ if key does
452 -- not exist in the map. If the key does exist, the function will
453 -- insert @f new_value old_value@.
454 --
455 -- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
456 -- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
457 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx"
458
459 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
460 insertWith f k x t
461 = insertWithKey (\_ x' y' -> f x' y') k x t
462
463 -- | Same as 'insertWith', but the combining function is applied strictly.
464 insertWith' :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
465 insertWith' f k x t
466 = insertWithKey' (\_ x' y' -> f x' y') k x t
467
468 -- | /O(min(n,W))/. Insert with a combining function.
469 -- @'insertWithKey' f key value mp@
470 -- will insert the pair (key, value) into @mp@ if key does
471 -- not exist in the map. If the key does exist, the function will
472 -- insert @f key new_value old_value@.
473 --
474 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
475 -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
476 -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
477 -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx"
478
479 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
480 insertWithKey f k x t = k `seq`
481 case t of
482 Bin p m l r
483 | nomatch k p m -> join k (Tip k x) p t
484 | zero k m -> Bin p m (insertWithKey f k x l) r
485 | otherwise -> Bin p m l (insertWithKey f k x r)
486 Tip ky y
487 | k==ky -> Tip k (f k x y)
488 | otherwise -> join k (Tip k x) ky t
489 Nil -> Tip k x
490
491 -- | Same as 'insertWithKey', but the combining function is applied strictly.
492 insertWithKey' :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
493 insertWithKey' f k x t = k `seq`
494 case t of
495 Bin p m l r
496 | nomatch k p m -> join k (Tip k x) p t
497 | zero k m -> Bin p m (insertWithKey' f k x l) r
498 | otherwise -> Bin p m l (insertWithKey' f k x r)
499 Tip ky y
500 | k==ky -> let x' = f k x y in seq x' (Tip k x')
501 | otherwise -> join k (Tip k x) ky t
502 Nil -> Tip k x
503
504 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
505 -- is a pair where the first element is equal to (@'lookup' k map@)
506 -- and the second element equal to (@'insertWithKey' f k x map@).
507 --
508 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
509 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
510 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
511 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
512 --
513 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
514 --
515 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
516 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
517 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
518
519 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
520 insertLookupWithKey f k x t = k `seq`
521 case t of
522 Bin p m l r
523 | nomatch k p m -> (Nothing,join k (Tip k x) p t)
524 | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
525 | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
526 Tip ky y
527 | k==ky -> (Just y,Tip k (f k x y))
528 | otherwise -> (Nothing,join k (Tip k x) ky t)
529 Nil -> (Nothing,Tip k x)
530
531
532 {--------------------------------------------------------------------
533 Deletion
534 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
535 --------------------------------------------------------------------}
536 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
537 -- a member of the map, the original map is returned.
538 --
539 -- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
540 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
541 -- > delete 5 empty == empty
542
543 delete :: Key -> IntMap a -> IntMap a
544 delete k t = k `seq`
545 case t of
546 Bin p m l r
547 | nomatch k p m -> t
548 | zero k m -> bin p m (delete k l) r
549 | otherwise -> bin p m l (delete k r)
550 Tip ky _
551 | k==ky -> Nil
552 | otherwise -> t
553 Nil -> Nil
554
555 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
556 -- a member of the map, the original map is returned.
557 --
558 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
559 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
560 -- > adjust ("new " ++) 7 empty == empty
561
562 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
563 adjust f k m
564 = adjustWithKey (\_ x -> f x) k m
565
566 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
567 -- a member of the map, the original map is returned.
568 --
569 -- > let f key x = (show key) ++ ":new " ++ x
570 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
571 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
572 -- > adjustWithKey f 7 empty == empty
573
574 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
575 adjustWithKey f
576 = updateWithKey (\k' x -> Just (f k' x))
577
578 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
579 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
580 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
581 --
582 -- > let f x = if x == "a" then Just "new a" else Nothing
583 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
584 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
585 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
586
587 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
588 update f
589 = updateWithKey (\_ x -> f x)
590
591 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
592 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
593 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
594 --
595 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
596 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
597 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
598 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
599
600 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
601 updateWithKey f k t = k `seq`
602 case t of
603 Bin p m l r
604 | nomatch k p m -> t
605 | zero k m -> bin p m (updateWithKey f k l) r
606 | otherwise -> bin p m l (updateWithKey f k r)
607 Tip ky y
608 | k==ky -> case (f k y) of
609 Just y' -> Tip ky y'
610 Nothing -> Nil
611 | otherwise -> t
612 Nil -> Nil
613
614 -- | /O(min(n,W))/. Lookup and update.
615 -- The function returns original value, if it is updated.
616 -- This is different behavior than 'Data.Map.updateLookupWithKey'.
617 -- Returns the original key value if the map entry is deleted.
618 --
619 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
620 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
621 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
622 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
623
624 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
625 updateLookupWithKey f k t = k `seq`
626 case t of
627 Bin p m l r
628 | nomatch k p m -> (Nothing,t)
629 | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
630 | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
631 Tip ky y
632 | k==ky -> case (f k y) of
633 Just y' -> (Just y,Tip ky y')
634 Nothing -> (Just y,Nil)
635 | otherwise -> (Nothing,t)
636 Nil -> (Nothing,Nil)
637
638
639
640 -- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
641 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
642 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
643 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
644 alter f k t = k `seq`
645 case t of
646 Bin p m l r
647 | nomatch k p m -> case f Nothing of
648 Nothing -> t
649 Just x -> join k (Tip k x) p t
650 | zero k m -> bin p m (alter f k l) r
651 | otherwise -> bin p m l (alter f k r)
652 Tip ky y
653 | k==ky -> case f (Just y) of
654 Just x -> Tip ky x
655 Nothing -> Nil
656 | otherwise -> case f Nothing of
657 Just x -> join k (Tip k x) ky t
658 Nothing -> Tip ky y
659 Nil -> case f Nothing of
660 Just x -> Tip k x
661 Nothing -> Nil
662
663
664 {--------------------------------------------------------------------
665 Union
666 --------------------------------------------------------------------}
667 -- | The union of a list of maps.
668 --
669 -- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
670 -- > == fromList [(3, "b"), (5, "a"), (7, "C")]
671 -- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])]
672 -- > == fromList [(3, "B3"), (5, "A3"), (7, "C")]
673
674 unions :: [IntMap a] -> IntMap a
675 unions xs
676 = foldlStrict union empty xs
677
678 -- | The union of a list of maps, with a combining operation.
679 --
680 -- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
681 -- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
682
683 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
684 unionsWith f ts
685 = foldlStrict (unionWith f) empty ts
686
687 -- | /O(n+m)/. The (left-biased) union of two maps.
688 -- It prefers the first map when duplicate keys are encountered,
689 -- i.e. (@'union' == 'unionWith' 'const'@).
690 --
691 -- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")]
692
693 union :: IntMap a -> IntMap a -> IntMap a
694 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
695 | shorter m1 m2 = union1
696 | shorter m2 m1 = union2
697 | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2)
698 | otherwise = join p1 t1 p2 t2
699 where
700 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
701 | zero p2 m1 = Bin p1 m1 (union l1 t2) r1
702 | otherwise = Bin p1 m1 l1 (union r1 t2)
703
704 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
705 | zero p1 m2 = Bin p2 m2 (union t1 l2) r2
706 | otherwise = Bin p2 m2 l2 (union t1 r2)
707
708 union (Tip k x) t = insert k x t
709 union t (Tip k x) = insertWith (\_ y -> y) k x t -- right bias
710 union Nil t = t
711 union t Nil = t
712
713 -- | /O(n+m)/. The union with a combining function.
714 --
715 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
716
717 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
718 unionWith f m1 m2
719 = unionWithKey (\_ x y -> f x y) m1 m2
720
721 -- | /O(n+m)/. The union with a combining function.
722 --
723 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
724 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
725
726 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
727 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
728 | shorter m1 m2 = union1
729 | shorter m2 m1 = union2
730 | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
731 | otherwise = join p1 t1 p2 t2
732 where
733 union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
734 | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1
735 | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2)
736
737 union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
738 | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2
739 | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2)
740
741 unionWithKey f (Tip k x) t = insertWithKey f k x t
742 unionWithKey f t (Tip k x) = insertWithKey (\k' x' y' -> f k' y' x') k x t -- right bias
743 unionWithKey _ Nil t = t
744 unionWithKey _ t Nil = t
745
746 {--------------------------------------------------------------------
747 Difference
748 --------------------------------------------------------------------}
749 -- | /O(n+m)/. Difference between two maps (based on keys).
750 --
751 -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b"
752
753 difference :: IntMap a -> IntMap b -> IntMap a
754 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
755 | shorter m1 m2 = difference1
756 | shorter m2 m1 = difference2
757 | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2)
758 | otherwise = t1
759 where
760 difference1 | nomatch p2 p1 m1 = t1
761 | zero p2 m1 = bin p1 m1 (difference l1 t2) r1
762 | otherwise = bin p1 m1 l1 (difference r1 t2)
763
764 difference2 | nomatch p1 p2 m2 = t1
765 | zero p1 m2 = difference t1 l2
766 | otherwise = difference t1 r2
767
768 difference t1@(Tip k _) t2
769 | member k t2 = Nil
770 | otherwise = t1
771
772 difference Nil _ = Nil
773 difference t (Tip k _) = delete k t
774 difference t Nil = t
775
776 -- | /O(n+m)/. Difference with a combining function.
777 --
778 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
779 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
780 -- > == singleton 3 "b:B"
781
782 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
783 differenceWith f m1 m2
784 = differenceWithKey (\_ x y -> f x y) m1 m2
785
786 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
787 -- encountered, the combining function is applied to the key and both values.
788 -- If it returns 'Nothing', the element is discarded (proper set difference).
789 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
790 --
791 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
792 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
793 -- > == singleton 3 "3:b|B"
794
795 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
796 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
797 | shorter m1 m2 = difference1
798 | shorter m2 m1 = difference2
799 | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
800 | otherwise = t1
801 where
802 difference1 | nomatch p2 p1 m1 = t1
803 | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1
804 | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2)
805
806 difference2 | nomatch p1 p2 m2 = t1
807 | zero p1 m2 = differenceWithKey f t1 l2
808 | otherwise = differenceWithKey f t1 r2
809
810 differenceWithKey f t1@(Tip k x) t2
811 = case lookup k t2 of
812 Just y -> case f k x y of
813 Just y' -> Tip k y'
814 Nothing -> Nil
815 Nothing -> t1
816
817 differenceWithKey _ Nil _ = Nil
818 differenceWithKey f t (Tip k y) = updateWithKey (\k' x -> f k' x y) k t
819 differenceWithKey _ t Nil = t
820
821
822 {--------------------------------------------------------------------
823 Intersection
824 --------------------------------------------------------------------}
825 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
826 --
827 -- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a"
828
829 intersection :: IntMap a -> IntMap b -> IntMap a
830 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
831 | shorter m1 m2 = intersection1
832 | shorter m2 m1 = intersection2
833 | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
834 | otherwise = Nil
835 where
836 intersection1 | nomatch p2 p1 m1 = Nil
837 | zero p2 m1 = intersection l1 t2
838 | otherwise = intersection r1 t2
839
840 intersection2 | nomatch p1 p2 m2 = Nil
841 | zero p1 m2 = intersection t1 l2
842 | otherwise = intersection t1 r2
843
844 intersection t1@(Tip k _) t2
845 | member k t2 = t1
846 | otherwise = Nil
847 intersection t (Tip k _)
848 = case lookup k t of
849 Just y -> Tip k y
850 Nothing -> Nil
851 intersection Nil _ = Nil
852 intersection _ Nil = Nil
853
854 -- | /O(n+m)/. The intersection with a combining function.
855 --
856 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
857
858 intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
859 intersectionWith f m1 m2
860 = intersectionWithKey (\_ x y -> f x y) m1 m2
861
862 -- | /O(n+m)/. The intersection with a combining function.
863 --
864 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
865 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
866
867 intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
868 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
869 | shorter m1 m2 = intersection1
870 | shorter m2 m1 = intersection2
871 | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
872 | otherwise = Nil
873 where
874 intersection1 | nomatch p2 p1 m1 = Nil
875 | zero p2 m1 = intersectionWithKey f l1 t2
876 | otherwise = intersectionWithKey f r1 t2
877
878 intersection2 | nomatch p1 p2 m2 = Nil
879 | zero p1 m2 = intersectionWithKey f t1 l2
880 | otherwise = intersectionWithKey f t1 r2
881
882 intersectionWithKey f (Tip k x) t2
883 = case lookup k t2 of
884 Just y -> Tip k (f k x y)
885 Nothing -> Nil
886 intersectionWithKey f t1 (Tip k y)
887 = case lookup k t1 of
888 Just x -> Tip k (f k x y)
889 Nothing -> Nil
890 intersectionWithKey _ Nil _ = Nil
891 intersectionWithKey _ _ Nil = Nil
892
893
894 {--------------------------------------------------------------------
895 Min\/Max
896 --------------------------------------------------------------------}
897
898 -- | /O(log n)/. Update the value at the minimal key.
899 --
900 -- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
901 -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
902
903 updateMinWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
904 updateMinWithKey f t
905 = case t of
906 Bin p m l r | m < 0 -> let t' = updateMinWithKeyUnsigned f r in Bin p m l t'
907 Bin p m l r -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r
908 Tip k y -> Tip k (f k y)
909 Nil -> error "maxView: empty map has no maximal element"
910
911 updateMinWithKeyUnsigned :: (Key -> a -> a) -> IntMap a -> IntMap a
912 updateMinWithKeyUnsigned f t
913 = case t of
914 Bin p m l r -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r
915 Tip k y -> Tip k (f k y)
916 Nil -> error "updateMinWithKeyUnsigned Nil"
917
918 -- | /O(log n)/. Update the value at the maximal key.
919 --
920 -- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
921 -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
922
923 updateMaxWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
924 updateMaxWithKey f t
925 = case t of
926 Bin p m l r | m < 0 -> let t' = updateMaxWithKeyUnsigned f l in Bin p m t' r
927 Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t'
928 Tip k y -> Tip k (f k y)
929 Nil -> error "maxView: empty map has no maximal element"
930
931 updateMaxWithKeyUnsigned :: (Key -> a -> a) -> IntMap a -> IntMap a
932 updateMaxWithKeyUnsigned f t
933 = case t of
934 Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t'
935 Tip k y -> Tip k (f k y)
936 Nil -> error "updateMaxWithKeyUnsigned Nil"
937
938
939 -- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and
940 -- the map stripped of that element, or 'Nothing' if passed an empty map.
941 --
942 -- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b")
943 -- > maxViewWithKey empty == Nothing
944
945 maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
946 maxViewWithKey t
947 = case t of
948 Bin p m l r | m < 0 -> let (result, t') = maxViewUnsigned l in Just (result, bin p m t' r)
949 Bin p m l r -> let (result, t') = maxViewUnsigned r in Just (result, bin p m l t')
950 Tip k y -> Just ((k,y), Nil)
951 Nil -> Nothing
952
953 maxViewUnsigned :: IntMap a -> ((Key, a), IntMap a)
954 maxViewUnsigned t
955 = case t of
956 Bin p m l r -> let (result,t') = maxViewUnsigned r in (result,bin p m l t')
957 Tip k y -> ((k,y), Nil)
958 Nil -> error "maxViewUnsigned Nil"
959
960 -- | /O(log n)/. Retrieves the minimal (key,value) pair of the map, and
961 -- the map stripped of that element, or 'Nothing' if passed an empty map.
962 --
963 -- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a")
964 -- > minViewWithKey empty == Nothing
965
966 minViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
967 minViewWithKey t
968 = case t of
969 Bin p m l r | m < 0 -> let (result, t') = minViewUnsigned r in Just (result, bin p m l t')
970 Bin p m l r -> let (result, t') = minViewUnsigned l in Just (result, bin p m t' r)
971 Tip k y -> Just ((k,y),Nil)
972 Nil -> Nothing
973
974 minViewUnsigned :: IntMap a -> ((Key, a), IntMap a)
975 minViewUnsigned t
976 = case t of
977 Bin p m l r -> let (result,t') = minViewUnsigned l in (result,bin p m t' r)
978 Tip k y -> ((k,y),Nil)
979 Nil -> error "minViewUnsigned Nil"
980
981
982 -- | /O(log n)/. Update the value at the maximal key.
983 --
984 -- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
985 -- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
986
987 updateMax :: (a -> a) -> IntMap a -> IntMap a
988 updateMax f = updateMaxWithKey (const f)
989
990 -- | /O(log n)/. Update the value at the minimal key.
991 --
992 -- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
993 -- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
994
995 updateMin :: (a -> a) -> IntMap a -> IntMap a
996 updateMin f = updateMinWithKey (const f)
997
998 -- Similar to the Arrow instance.
999 first :: (a -> c) -> (a, b) -> (c, b)
1000 first f (x,y) = (f x,y)
1001
1002 -- | /O(log n)/. Retrieves the maximal key of the map, and the map
1003 -- stripped of that element, or 'Nothing' if passed an empty map.
1004 maxView :: IntMap a -> Maybe (a, IntMap a)
1005 maxView t = liftM (first snd) (maxViewWithKey t)
1006
1007 -- | /O(log n)/. Retrieves the minimal key of the map, and the map
1008 -- stripped of that element, or 'Nothing' if passed an empty map.
1009 minView :: IntMap a -> Maybe (a, IntMap a)
1010 minView t = liftM (first snd) (minViewWithKey t)
1011
1012 -- | /O(log n)/. Delete and find the maximal element.
1013 deleteFindMax :: IntMap a -> (a, IntMap a)
1014 deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxView
1015
1016 -- | /O(log n)/. Delete and find the minimal element.
1017 deleteFindMin :: IntMap a -> (a, IntMap a)
1018 deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minView
1019
1020 -- | /O(log n)/. The minimal key of the map.
1021 findMin :: IntMap a -> (Key, a)
1022 findMin Nil = error $ "findMin: empty map has no minimal element"
1023 findMin (Tip k v) = (k,v)
1024 findMin (Bin _ m l r)
1025 | m < 0 = go r
1026 | otherwise = go l
1027 where go (Tip k v) = (k,v)
1028 go (Bin _ _ l' _) = go l'
1029 go Nil = error "findMax Nil"
1030
1031 -- | /O(log n)/. The maximal key of the map.
1032 findMax :: IntMap a -> (Key, a)
1033 findMax Nil = error $ "findMax: empty map has no maximal element"
1034 findMax (Tip k v) = (k,v)
1035 findMax (Bin _ m l r)
1036 | m < 0 = go l
1037 | otherwise = go r
1038 where go (Tip k v) = (k,v)
1039 go (Bin _ _ _ r') = go r'
1040 go Nil = error "findMax Nil"
1041
1042 -- | /O(log n)/. Delete the minimal key. An error is thrown if the IntMap is already empty.
1043 -- Note, this is not the same behavior Map.
1044 deleteMin :: IntMap a -> IntMap a
1045 deleteMin = maybe (error "deleteMin: empty map has no minimal element") snd . minView
1046
1047 -- | /O(log n)/. Delete the maximal key. An error is thrown if the IntMap is already empty.
1048 -- Note, this is not the same behavior Map.
1049 deleteMax :: IntMap a -> IntMap a
1050 deleteMax = maybe (error "deleteMax: empty map has no maximal element") snd . maxView
1051
1052
1053 {--------------------------------------------------------------------
1054 Submap
1055 --------------------------------------------------------------------}
1056 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1057 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
1058 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1059 isProperSubmapOf m1 m2
1060 = isProperSubmapOfBy (==) m1 m2
1061
1062 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
1063 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
1064 @m1@ and @m2@ are not equal,
1065 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1066 applied to their respective values. For example, the following
1067 expressions are all 'True':
1068
1069 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1070 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1071
1072 But the following are all 'False':
1073
1074 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
1075 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1076 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1077 -}
1078 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1079 isProperSubmapOfBy predicate t1 t2
1080 = case submapCmp predicate t1 t2 of
1081 LT -> True
1082 _ -> False
1083
1084 submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering
1085 submapCmp predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1086 | shorter m1 m2 = GT
1087 | shorter m2 m1 = submapCmpLt
1088 | p1 == p2 = submapCmpEq
1089 | otherwise = GT -- disjoint
1090 where
1091 submapCmpLt | nomatch p1 p2 m2 = GT
1092 | zero p1 m2 = submapCmp predicate t1 l2
1093 | otherwise = submapCmp predicate t1 r2
1094 submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of
1095 (GT,_ ) -> GT
1096 (_ ,GT) -> GT
1097 (EQ,EQ) -> EQ
1098 _ -> LT
1099
1100 submapCmp _ (Bin _ _ _ _) _ = GT
1101 submapCmp predicate (Tip kx x) (Tip ky y)
1102 | (kx == ky) && predicate x y = EQ
1103 | otherwise = GT -- disjoint
1104 submapCmp predicate (Tip k x) t
1105 = case lookup k t of
1106 Just y | predicate x y -> LT
1107 _ -> GT -- disjoint
1108 submapCmp _ Nil Nil = EQ
1109 submapCmp _ Nil _ = LT
1110
1111 -- | /O(n+m)/. Is this a submap?
1112 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
1113 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
1114 isSubmapOf m1 m2
1115 = isSubmapOfBy (==) m1 m2
1116
1117 {- | /O(n+m)/.
1118 The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
1119 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
1120 applied to their respective values. For example, the following
1121 expressions are all 'True':
1122
1123 > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1124 > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1125 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
1126
1127 But the following are all 'False':
1128
1129 > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
1130 > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
1131 > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
1132 -}
1133 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
1134 isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1135 | shorter m1 m2 = False
1136 | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy predicate t1 l2
1137 else isSubmapOfBy predicate t1 r2)
1138 | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2
1139 isSubmapOfBy _ (Bin _ _ _ _) _ = False
1140 isSubmapOfBy predicate (Tip k x) t = case lookup k t of
1141 Just y -> predicate x y
1142 Nothing -> False
1143 isSubmapOfBy _ Nil _ = True
1144
1145 {--------------------------------------------------------------------
1146 Mapping
1147 --------------------------------------------------------------------}
1148 -- | /O(n)/. Map a function over all values in the map.
1149 --
1150 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
1151
1152 map :: (a -> b) -> IntMap a -> IntMap b
1153 map f = mapWithKey (\_ x -> f x)
1154
1155 -- | /O(n)/. Map a function over all values in the map.
1156 --
1157 -- > let f key x = (show key) ++ ":" ++ x
1158 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
1159
1160 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
1161 mapWithKey f t
1162 = case t of
1163 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
1164 Tip k x -> Tip k (f k x)
1165 Nil -> Nil
1166
1167 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
1168 -- argument through the map in ascending order of keys.
1169 --
1170 -- > let f a b = (a ++ b, b ++ "X")
1171 -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")])
1172
1173 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1174 mapAccum f = mapAccumWithKey (\a' _ x -> f a' x)
1175
1176 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
1177 -- argument through the map in ascending order of keys.
1178 --
1179 -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X")
1180 -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")])
1181
1182 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1183 mapAccumWithKey f a t
1184 = mapAccumL f a t
1185
1186 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
1187 -- argument through the map in ascending order of keys.
1188 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1189 mapAccumL f a t
1190 = case t of
1191 Bin p m l r -> let (a1,l') = mapAccumL f a l
1192 (a2,r') = mapAccumL f a1 r
1193 in (a2,Bin p m l' r')
1194 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1195 Nil -> (a,Nil)
1196
1197 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
1198 -- argument through the map in descending order of keys.
1199 mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
1200 mapAccumRWithKey f a t
1201 = case t of
1202 Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r
1203 (a2,l') = mapAccumRWithKey f a1 l
1204 in (a2,Bin p m l' r')
1205 Tip k x -> let (a',x') = f a k x in (a',Tip k x')
1206 Nil -> (a,Nil)
1207
1208 {--------------------------------------------------------------------
1209 Filter
1210 --------------------------------------------------------------------}
1211 -- | /O(n)/. Filter all values that satisfy some predicate.
1212 --
1213 -- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
1214 -- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty
1215 -- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty
1216
1217 filter :: (a -> Bool) -> IntMap a -> IntMap a
1218 filter p m
1219 = filterWithKey (\_ x -> p x) m
1220
1221 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
1222 --
1223 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
1224
1225 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
1226 filterWithKey predicate t
1227 = case t of
1228 Bin p m l r
1229 -> bin p m (filterWithKey predicate l) (filterWithKey predicate r)
1230 Tip k x
1231 | predicate k x -> t
1232 | otherwise -> Nil
1233 Nil -> Nil
1234
1235 -- | /O(n)/. Partition the map according to some predicate. The first
1236 -- map contains all elements that satisfy the predicate, the second all
1237 -- elements that fail the predicate. See also 'split'.
1238 --
1239 -- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1240 -- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1241 -- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1242
1243 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1244 partition p m
1245 = partitionWithKey (\_ x -> p x) m
1246
1247 -- | /O(n)/. Partition the map according to some predicate. The first
1248 -- map contains all elements that satisfy the predicate, the second all
1249 -- elements that fail the predicate. See also 'split'.
1250 --
1251 -- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b")
1252 -- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
1253 -- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
1254
1255 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
1256 partitionWithKey predicate t
1257 = case t of
1258 Bin p m l r
1259 -> let (l1,l2) = partitionWithKey predicate l
1260 (r1,r2) = partitionWithKey predicate r
1261 in (bin p m l1 r1, bin p m l2 r2)
1262 Tip k x
1263 | predicate k x -> (t,Nil)
1264 | otherwise -> (Nil,t)
1265 Nil -> (Nil,Nil)
1266
1267 -- | /O(n)/. Map values and collect the 'Just' results.
1268 --
1269 -- > let f x = if x == "a" then Just "new a" else Nothing
1270 -- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
1271
1272 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
1273 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
1274
1275 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
1276 --
1277 -- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
1278 -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
1279
1280 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
1281 mapMaybeWithKey f (Bin p m l r)
1282 = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
1283 mapMaybeWithKey f (Tip k x) = case f k x of
1284 Just y -> Tip k y
1285 Nothing -> Nil
1286 mapMaybeWithKey _ Nil = Nil
1287
1288 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
1289 --
1290 -- > let f a = if a < "c" then Left a else Right a
1291 -- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1292 -- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
1293 -- >
1294 -- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1295 -- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1296
1297 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1298 mapEither f m
1299 = mapEitherWithKey (\_ x -> f x) m
1300
1301 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
1302 --
1303 -- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
1304 -- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1305 -- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
1306 -- >
1307 -- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1308 -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
1309
1310 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1311 mapEitherWithKey f (Bin p m l r)
1312 = (bin p m l1 r1, bin p m l2 r2)
1313 where
1314 (l1,l2) = mapEitherWithKey f l
1315 (r1,r2) = mapEitherWithKey f r
1316 mapEitherWithKey f (Tip k x) = case f k x of
1317 Left y -> (Tip k y, Nil)
1318 Right z -> (Nil, Tip k z)
1319 mapEitherWithKey _ Nil = (Nil, Nil)
1320
1321 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@
1322 -- where all keys in @map1@ are lower than @k@ and all keys in
1323 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
1324 --
1325 -- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")])
1326 -- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a")
1327 -- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
1328 -- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty)
1329 -- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty)
1330
1331 split :: Key -> IntMap a -> (IntMap a,IntMap a)
1332 split k t
1333 = case t of
1334 Bin _ m l r
1335 | m < 0 -> (if k >= 0 -- handle negative numbers.
1336 then let (lt,gt) = split' k l in (union r lt, gt)
1337 else let (lt,gt) = split' k r in (lt, union gt l))
1338 | otherwise -> split' k t
1339 Tip ky _
1340 | k>ky -> (t,Nil)
1341 | k<ky -> (Nil,t)
1342 | otherwise -> (Nil,Nil)
1343 Nil -> (Nil,Nil)
1344
1345 split' :: Key -> IntMap a -> (IntMap a,IntMap a)
1346 split' k t
1347 = case t of
1348 Bin p m l r
1349 | nomatch k p m -> if k>p then (t,Nil) else (Nil,t)
1350 | zero k m -> let (lt,gt) = split k l in (lt,union gt r)
1351 | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
1352 Tip ky _
1353 | k>ky -> (t,Nil)
1354 | k<ky -> (Nil,t)
1355 | otherwise -> (Nil,Nil)
1356 Nil -> (Nil,Nil)
1357
1358 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
1359 -- key was found in the original map.
1360 --
1361 -- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")])
1362 -- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a")
1363 -- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a")
1364 -- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty)
1365 -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty)
1366
1367 splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
1368 splitLookup k t
1369 = case t of
1370 Bin _ m l r
1371 | m < 0 -> (if k >= 0 -- handle negative numbers.
1372 then let (lt,found,gt) = splitLookup' k l in (union r lt,found, gt)
1373 else let (lt,found,gt) = splitLookup' k r in (lt,found, union gt l))
1374 | otherwise -> splitLookup' k t
1375 Tip ky y
1376 | k>ky -> (t,Nothing,Nil)
1377 | k<ky -> (Nil,Nothing,t)
1378 | otherwise -> (Nil,Just y,Nil)
1379 Nil -> (Nil,Nothing,Nil)
1380
1381 splitLookup' :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
1382 splitLookup' k t
1383 = case t of
1384 Bin p m l r
1385 | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t)
1386 | zero k m -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
1387 | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
1388 Tip ky y
1389 | k>ky -> (t,Nothing,Nil)
1390 | k<ky -> (Nil,Nothing,t)
1391 | otherwise -> (Nil,Just y,Nil)
1392 Nil -> (Nil,Nothing,Nil)
1393
1394 {--------------------------------------------------------------------
1395 Fold
1396 --------------------------------------------------------------------}
1397 -- | /O(n)/. Fold the values in the map using the given right-associative
1398 -- binary operator. This function is an equivalent of 'foldr' and is present
1399 -- for compatibility only.
1400 --
1401 -- /Please note that fold will be deprecated in the future and removed./
1402 fold :: (a -> b -> b) -> b -> IntMap a -> b
1403 fold = foldr
1404 {-# INLINE fold #-}
1405
1406 -- | /O(n)/. Fold the values in the map using the given right-associative
1407 -- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@.
1408 --
1409 -- For example,
1410 --
1411 -- > elems map = foldr (:) [] map
1412 --
1413 -- > let f a len = len + (length a)
1414 -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1415 foldr :: (a -> b -> b) -> b -> IntMap a -> b
1416 foldr f z t =
1417 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1418 _ -> go z t
1419 where
1420 go z' Nil = z'
1421 go z' (Tip _ x) = f x z'
1422 go z' (Bin _ _ l r) = go (go z' r) l
1423 {-# INLINE foldr #-}
1424
1425 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
1426 -- evaluated before using the result in the next application. This
1427 -- function is strict in the starting value.
1428 foldr' :: (a -> b -> b) -> b -> IntMap a -> b
1429 foldr' f z t =
1430 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1431 _ -> go z t
1432 where
1433 STRICT_1_OF_2(go)
1434 go z' Nil = z'
1435 go z' (Tip _ x) = f x z'
1436 go z' (Bin _ _ l r) = go (go z' r) l
1437 {-# INLINE foldr' #-}
1438
1439 -- | /O(n)/. Fold the values in the map using the given left-associative
1440 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
1441 --
1442 -- For example,
1443 --
1444 -- > elems = reverse . foldl (flip (:)) []
1445 --
1446 -- > let f len a = len + (length a)
1447 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1448 foldl :: (a -> b -> a) -> a -> IntMap b -> a
1449 foldl f z t =
1450 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1451 _ -> go z t
1452 where
1453 go z' Nil = z'
1454 go z' (Tip _ x) = f z' x
1455 go z' (Bin _ _ l r) = go (go z' l) r
1456 {-# INLINE foldl #-}
1457
1458 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
1459 -- evaluated before using the result in the next application. This
1460 -- function is strict in the starting value.
1461 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
1462 foldl' f z t =
1463 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1464 _ -> go z t
1465 where
1466 STRICT_1_OF_2(go)
1467 go z' Nil = z'
1468 go z' (Tip _ x) = f z' x
1469 go z' (Bin _ _ l r) = go (go z' l) r
1470 {-# INLINE foldl' #-}
1471
1472 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1473 -- binary operator. This function is an equivalent of 'foldrWithKey' and is present
1474 -- for compatibility only.
1475 --
1476 -- /Please note that foldWithKey will be deprecated in the future and removed./
1477 foldWithKey :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1478 foldWithKey = foldrWithKey
1479 {-# INLINE foldWithKey #-}
1480
1481 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1482 -- binary operator, such that
1483 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
1484 --
1485 -- For example,
1486 --
1487 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
1488 --
1489 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1490 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
1491 foldrWithKey :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1492 foldrWithKey f z t =
1493 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1494 _ -> go z t
1495 where
1496 go z' Nil = z'
1497 go z' (Tip kx x) = f kx x z'
1498 go z' (Bin _ _ l r) = go (go z' r) l
1499 {-# INLINE foldrWithKey #-}
1500
1501 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
1502 -- evaluated before using the result in the next application. This
1503 -- function is strict in the starting value.
1504 foldrWithKey' :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1505 foldrWithKey' f z t =
1506 case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before
1507 _ -> go z t
1508 where
1509 STRICT_1_OF_2(go)
1510 go z' Nil = z'
1511 go z' (Tip kx x) = f kx x z'
1512 go z' (Bin _ _ l r) = go (go z' r) l
1513 {-# INLINE foldrWithKey' #-}
1514
1515 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
1516 -- binary operator, such that
1517 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
1518 --
1519 -- For example,
1520 --
1521 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
1522 --
1523 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1524 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
1525 foldlWithKey :: (a -> Int -> b -> a) -> a -> IntMap b -> a
1526 foldlWithKey f z t =
1527 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1528 _ -> go z t
1529 where
1530 go z' Nil = z'
1531 go z' (Tip kx x) = f z' kx x
1532 go z' (Bin _ _ l r) = go (go z' l) r
1533 {-# INLINE foldlWithKey #-}
1534
1535 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
1536 -- evaluated before using the result in the next application. This
1537 -- function is strict in the starting value.
1538 foldlWithKey' :: (a -> Int -> b -> a) -> a -> IntMap b -> a
1539 foldlWithKey' f z t =
1540 case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before
1541 _ -> go z t
1542 where
1543 STRICT_1_OF_2(go)
1544 go z' Nil = z'
1545 go z' (Tip kx x) = f z' kx x
1546 go z' (Bin _ _ l r) = go (go z' l) r
1547 {-# INLINE foldlWithKey' #-}
1548
1549 {--------------------------------------------------------------------
1550 List variations
1551 --------------------------------------------------------------------}
1552 -- | /O(n)/.
1553 -- Return all elements of the map in the ascending order of their keys.
1554 --
1555 -- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
1556 -- > elems empty == []
1557
1558 elems :: IntMap a -> [a]
1559 elems
1560 = foldr (:) []
1561
1562 -- | /O(n)/. Return all keys of the map in ascending order.
1563 --
1564 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
1565 -- > keys empty == []
1566
1567 keys :: IntMap a -> [Key]
1568 keys
1569 = foldrWithKey (\k _ ks -> k:ks) []
1570
1571 -- | /O(n*min(n,W))/. The set of all keys of the map.
1572 --
1573 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
1574 -- > keysSet empty == Data.IntSet.empty
1575
1576 keysSet :: IntMap a -> IntSet.IntSet
1577 keysSet m = IntSet.fromDistinctAscList (keys m)
1578
1579
1580 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
1581 --
1582 -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1583 -- > assocs empty == []
1584
1585 assocs :: IntMap a -> [(Key,a)]
1586 assocs m
1587 = toList m
1588
1589
1590 {--------------------------------------------------------------------
1591 Lists
1592 --------------------------------------------------------------------}
1593 -- | /O(n)/. Convert the map to a list of key\/value pairs.
1594 --
1595 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1596 -- > toList empty == []
1597
1598 toList :: IntMap a -> [(Key,a)]
1599 toList
1600 = foldrWithKey (\k x xs -> (k,x):xs) []
1601
1602 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
1603 -- keys are in ascending order.
1604 --
1605 -- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1606
1607 toAscList :: IntMap a -> [(Key,a)]
1608 toAscList t
1609 = -- NOTE: the following algorithm only works for big-endian trees
1610 let (pos,neg) = span (\(k,_) -> k >=0) (foldrWithKey (\k x xs -> (k,x):xs) [] t) in neg ++ pos
1611
1612 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
1613 --
1614 -- > fromList [] == empty
1615 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
1616 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
1617
1618 fromList :: [(Key,a)] -> IntMap a
1619 fromList xs
1620 = foldlStrict ins empty xs
1621 where
1622 ins t (k,x) = insert k x t
1623
1624 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1625 --
1626 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
1627 -- > fromListWith (++) [] == empty
1628
1629 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1630 fromListWith f xs
1631 = fromListWithKey (\_ x y -> f x y) xs
1632
1633 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
1634 --
1635 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
1636 -- > fromListWith (++) [] == empty
1637
1638 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1639 fromListWithKey f xs
1640 = foldlStrict ins empty xs
1641 where
1642 ins t (k,x) = insertWithKey f k x t
1643
1644 -- | /O(n)/. Build a map from a list of key\/value pairs where
1645 -- the keys are in ascending order.
1646 --
1647 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1648 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
1649
1650 fromAscList :: [(Key,a)] -> IntMap a
1651 fromAscList xs
1652 = fromAscListWithKey (\_ x _ -> x) xs
1653
1654 -- | /O(n)/. Build a map from a list of key\/value pairs where
1655 -- the keys are in ascending order, with a combining function on equal keys.
1656 -- /The precondition (input list is ascending) is not checked./
1657 --
1658 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1659
1660 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1661 fromAscListWith f xs
1662 = fromAscListWithKey (\_ x y -> f x y) xs
1663
1664 -- | /O(n)/. Build a map from a list of key\/value pairs where
1665 -- the keys are in ascending order, with a combining function on equal keys.
1666 -- /The precondition (input list is ascending) is not checked./
1667 --
1668 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1669
1670 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1671 fromAscListWithKey _ [] = Nil
1672 fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
1673 where
1674 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1675 combineEq z [] = [z]
1676 combineEq z@(kz,zz) (x@(kx,xx):xs)
1677 | kx==kz = let yy = f kx xx zz in combineEq (kx,yy) xs
1678 | otherwise = z:combineEq x xs
1679
1680 -- | /O(n)/. Build a map from a list of key\/value pairs where
1681 -- the keys are in ascending order and all distinct.
1682 -- /The precondition (input list is strictly ascending) is not checked./
1683 --
1684 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1685
1686 #ifdef __GLASGOW_HASKELL__
1687 fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a
1688 #else
1689 fromDistinctAscList :: [(Key,a)] -> IntMap a
1690 #endif
1691 fromDistinctAscList [] = Nil
1692 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
1693 where
1694 work (kx,vx) [] stk = finish kx (Tip kx vx) stk
1695 work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
1696
1697 #ifdef __GLASGOW_HASKELL__
1698 reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
1699 #endif
1700 reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
1701 reduce z zs m px tx stk@(Push py ty stk') =
1702 let mxy = branchMask px py
1703 pxy = mask px mxy
1704 in if shorter m mxy
1705 then reduce z zs m pxy (Bin pxy mxy ty tx) stk'
1706 else work z zs (Push px tx stk)
1707
1708 finish _ t Nada = t
1709 finish px tx (Push py ty stk) = finish p (join py ty px tx) stk
1710 where m = branchMask px py
1711 p = mask px m
1712
1713 data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada
1714
1715
1716 {--------------------------------------------------------------------
1717 Eq
1718 --------------------------------------------------------------------}
1719 instance Eq a => Eq (IntMap a) where
1720 t1 == t2 = equal t1 t2
1721 t1 /= t2 = nequal t1 t2
1722
1723 equal :: Eq a => IntMap a -> IntMap a -> Bool
1724 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1725 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
1726 equal (Tip kx x) (Tip ky y)
1727 = (kx == ky) && (x==y)
1728 equal Nil Nil = True
1729 equal _ _ = False
1730
1731 nequal :: Eq a => IntMap a -> IntMap a -> Bool
1732 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1733 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
1734 nequal (Tip kx x) (Tip ky y)
1735 = (kx /= ky) || (x/=y)
1736 nequal Nil Nil = False
1737 nequal _ _ = True
1738
1739 {--------------------------------------------------------------------
1740 Ord
1741 --------------------------------------------------------------------}
1742
1743 instance Ord a => Ord (IntMap a) where
1744 compare m1 m2 = compare (toList m1) (toList m2)
1745
1746 {--------------------------------------------------------------------
1747 Functor
1748 --------------------------------------------------------------------}
1749
1750 instance Functor IntMap where
1751 fmap = map
1752
1753 {--------------------------------------------------------------------
1754 Show
1755 --------------------------------------------------------------------}
1756
1757 instance Show a => Show (IntMap a) where
1758 showsPrec d m = showParen (d > 10) $
1759 showString "fromList " . shows (toList m)
1760
1761 {-
1762 XXX unused code
1763
1764 showMap :: (Show a) => [(Key,a)] -> ShowS
1765 showMap []
1766 = showString "{}"
1767 showMap (x:xs)
1768 = showChar '{' . showElem x . showTail xs
1769 where
1770 showTail [] = showChar '}'
1771 showTail (x':xs') = showChar ',' . showElem x' . showTail xs'
1772
1773 showElem (k,v) = shows k . showString ":=" . shows v
1774 -}
1775
1776 {--------------------------------------------------------------------
1777 Read
1778 --------------------------------------------------------------------}
1779 instance (Read e) => Read (IntMap e) where
1780 #ifdef __GLASGOW_HASKELL__
1781 readPrec = parens $ prec 10 $ do
1782 Ident "fromList" <- lexP
1783 xs <- readPrec
1784 return (fromList xs)
1785
1786 readListPrec = readListPrecDefault
1787 #else
1788 readsPrec p = readParen (p > 10) $ \ r -> do
1789 ("fromList",s) <- lex r
1790 (xs,t) <- reads s
1791 return (fromList xs,t)
1792 #endif
1793
1794 {--------------------------------------------------------------------
1795 Typeable
1796 --------------------------------------------------------------------}
1797
1798 #include "Typeable.h"
1799 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1800
1801 {--------------------------------------------------------------------
1802 Debugging
1803 --------------------------------------------------------------------}
1804 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1805 -- in a compressed, hanging format.
1806 showTree :: Show a => IntMap a -> String
1807 showTree s
1808 = showTreeWith True False s
1809
1810
1811 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1812 the tree that implements the map. If @hang@ is
1813 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1814 @wide@ is 'True', an extra wide version is shown.
1815 -}
1816 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1817 showTreeWith hang wide t
1818 | hang = (showsTreeHang wide [] t) ""
1819 | otherwise = (showsTree wide [] [] t) ""
1820
1821 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1822 showsTree wide lbars rbars t
1823 = case t of
1824 Bin p m l r
1825 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1826 showWide wide rbars .
1827 showsBars lbars . showString (showBin p m) . showString "\n" .
1828 showWide wide lbars .
1829 showsTree wide (withEmpty lbars) (withBar lbars) l
1830 Tip k x
1831 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1832 Nil -> showsBars lbars . showString "|\n"
1833
1834 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1835 showsTreeHang wide bars t
1836 = case t of
1837 Bin p m l r
1838 -> showsBars bars . showString (showBin p m) . showString "\n" .
1839 showWide wide bars .
1840 showsTreeHang wide (withBar bars) l .
1841 showWide wide bars .
1842 showsTreeHang wide (withEmpty bars) r
1843 Tip k x
1844 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1845 Nil -> showsBars bars . showString "|\n"
1846
1847 showBin :: Prefix -> Mask -> String
1848 showBin _ _
1849 = "*" -- ++ show (p,m)
1850
1851 showWide :: Bool -> [String] -> String -> String
1852 showWide wide bars
1853 | wide = showString (concat (reverse bars)) . showString "|\n"
1854 | otherwise = id
1855
1856 showsBars :: [String] -> ShowS
1857 showsBars bars
1858 = case bars of
1859 [] -> id
1860 _ -> showString (concat (reverse (tail bars))) . showString node
1861
1862 node :: String
1863 node = "+--"
1864
1865 withBar, withEmpty :: [String] -> [String]
1866 withBar bars = "| ":bars
1867 withEmpty bars = " ":bars
1868
1869
1870 {--------------------------------------------------------------------
1871 Helpers
1872 --------------------------------------------------------------------}
1873 {--------------------------------------------------------------------
1874 Join
1875 --------------------------------------------------------------------}
1876 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1877 join p1 t1 p2 t2
1878 | zero p1 m = Bin p m t1 t2
1879 | otherwise = Bin p m t2 t1
1880 where
1881 m = branchMask p1 p2
1882 p = mask p1 m
1883 {-# INLINE join #-}
1884
1885 {--------------------------------------------------------------------
1886 @bin@ assures that we never have empty trees within a tree.
1887 --------------------------------------------------------------------}
1888 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1889 bin _ _ l Nil = l
1890 bin _ _ Nil r = r
1891 bin p m l r = Bin p m l r
1892 {-# INLINE bin #-}
1893
1894
1895 {--------------------------------------------------------------------
1896 Endian independent bit twiddling
1897 --------------------------------------------------------------------}
1898 zero :: Key -> Mask -> Bool
1899 zero i m
1900 = (natFromInt i) .&. (natFromInt m) == 0
1901 {-# INLINE zero #-}
1902
1903 nomatch,match :: Key -> Prefix -> Mask -> Bool
1904 nomatch i p m
1905 = (mask i m) /= p
1906 {-# INLINE nomatch #-}
1907
1908 match i p m
1909 = (mask i m) == p
1910 {-# INLINE match #-}
1911
1912 mask :: Key -> Mask -> Prefix
1913 mask i m
1914 = maskW (natFromInt i) (natFromInt m)
1915 {-# INLINE mask #-}
1916
1917
1918 {--------------------------------------------------------------------
1919 Big endian operations
1920 --------------------------------------------------------------------}
1921 maskW :: Nat -> Nat -> Prefix
1922 maskW i m
1923 = intFromNat (i .&. (complement (m-1) `xor` m))
1924 {-# INLINE maskW #-}
1925
1926 shorter :: Mask -> Mask -> Bool
1927 shorter m1 m2
1928 = (natFromInt m1) > (natFromInt m2)
1929 {-# INLINE shorter #-}
1930
1931 branchMask :: Prefix -> Prefix -> Mask
1932 branchMask p1 p2
1933 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1934 {-# INLINE branchMask #-}
1935
1936 {----------------------------------------------------------------------
1937 Finding the highest bit (mask) in a word [x] can be done efficiently in
1938 three ways:
1939 * convert to a floating point value and the mantissa tells us the
1940 [log2(x)] that corresponds with the highest bit position. The mantissa
1941 is retrieved either via the standard C function [frexp] or by some bit
1942 twiddling on IEEE compatible numbers (float). Note that one needs to
1943 use at least [double] precision for an accurate mantissa of 32 bit
1944 numbers.
1945 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1946 * use processor specific assembler instruction (asm).
1947
1948 The most portable way would be [bit], but is it efficient enough?
1949 I have measured the cycle counts of the different methods on an AMD
1950 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1951
1952 highestBitMask: method cycles
1953 --------------
1954 frexp 200
1955 float 33
1956 bit 11
1957 asm 12
1958
1959 highestBit: method cycles
1960 --------------
1961 frexp 195
1962 float 33
1963 bit 11
1964 asm 11
1965
1966 Wow, the bit twiddling is on today's RISC like machines even faster
1967 than a single CISC instruction (BSR)!
1968 ----------------------------------------------------------------------}
1969
1970 {----------------------------------------------------------------------
1971 [highestBitMask] returns a word where only the highest bit is set.
1972 It is found by first setting all bits in lower positions than the
1973 highest bit and than taking an exclusive or with the original value.
1974 Allthough the function may look expensive, GHC compiles this into
1975 excellent C code that subsequently compiled into highly efficient
1976 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1977 ----------------------------------------------------------------------}
1978 highestBitMask :: Nat -> Nat
1979 highestBitMask x0
1980 = case (x0 .|. shiftRL x0 1) of
1981 x1 -> case (x1 .|. shiftRL x1 2) of
1982 x2 -> case (x2 .|. shiftRL x2 4) of
1983 x3 -> case (x3 .|. shiftRL x3 8) of
1984 x4 -> case (x4 .|. shiftRL x4 16) of
1985 x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms
1986 x6 -> (x6 `xor` (shiftRL x6 1))
1987 {-# INLINE highestBitMask #-}
1988
1989
1990 {--------------------------------------------------------------------
1991 Utilities
1992 --------------------------------------------------------------------}
1993
1994 foldlStrict :: (a -> b -> a) -> a -> [b] -> a
1995 foldlStrict f = go
1996 where
1997 go z [] = z
1998 go z (x:xs) = let z' = f z x in z' `seq` go z' xs
1999 {-# INLINE foldlStrict #-}