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