Improve Int{Set,Map}.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 -> -- Use lambda t to be inlinable with two arguments only.
1387 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1388 | otherwise -> go (go z r) l
1389 _ -> go z t
1390 where
1391 go z' Nil = z'
1392 go z' (Tip _ x) = f x z'
1393 go z' (Bin _ _ l r) = go (go z' r) l
1394 {-# INLINE foldr #-}
1395
1396 -- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
1397 -- evaluated before using the result in the next application. This
1398 -- function is strict in the starting value.
1399 foldr' :: (a -> b -> b) -> b -> IntMap a -> b
1400 foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1401 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1402 | otherwise -> go (go z r) l
1403 _ -> go z t
1404 where
1405 STRICT_1_OF_2(go)
1406 go z' Nil = z'
1407 go z' (Tip _ x) = f x z'
1408 go z' (Bin _ _ l r) = go (go z' r) l
1409 {-# INLINE foldr' #-}
1410
1411 -- | /O(n)/. Fold the values in the map using the given left-associative
1412 -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@.
1413 --
1414 -- For example,
1415 --
1416 -- > elems = reverse . foldl (flip (:)) []
1417 --
1418 -- > let f len a = len + (length a)
1419 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
1420 foldl :: (a -> b -> a) -> a -> IntMap b -> a
1421 foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1422 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1423 | otherwise -> go (go z l) r
1424 _ -> go z t
1425 where
1426 go z' Nil = z'
1427 go z' (Tip _ x) = f z' x
1428 go z' (Bin _ _ l r) = go (go z' l) r
1429 {-# INLINE foldl #-}
1430
1431 -- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
1432 -- evaluated before using the result in the next application. This
1433 -- function is strict in the starting value.
1434 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
1435 foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1436 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1437 | otherwise -> go (go z l) r
1438 _ -> go z t
1439 where
1440 STRICT_1_OF_2(go)
1441 go z' Nil = z'
1442 go z' (Tip _ x) = f z' x
1443 go z' (Bin _ _ l r) = go (go z' l) r
1444 {-# INLINE foldl' #-}
1445
1446 -- | /O(n)/. Fold the keys and values in the map using the given right-associative
1447 -- binary operator, such that
1448 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
1449 --
1450 -- For example,
1451 --
1452 -- > keys map = foldrWithKey (\k x ks -> k:ks) [] map
1453 --
1454 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1455 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
1456 foldrWithKey :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1457 foldrWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1458 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1459 | otherwise -> go (go z r) l
1460 _ -> go z t
1461 where
1462 go z' Nil = z'
1463 go z' (Tip kx x) = f kx x z'
1464 go z' (Bin _ _ l r) = go (go z' r) l
1465 {-# INLINE foldrWithKey #-}
1466
1467 -- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is
1468 -- evaluated before using the result in the next application. This
1469 -- function is strict in the starting value.
1470 foldrWithKey' :: (Int -> a -> b -> b) -> b -> IntMap a -> b
1471 foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1472 case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before
1473 | otherwise -> go (go z r) l
1474 _ -> go z t
1475 where
1476 STRICT_1_OF_2(go)
1477 go z' Nil = z'
1478 go z' (Tip kx x) = f kx x z'
1479 go z' (Bin _ _ l r) = go (go z' r) l
1480 {-# INLINE foldrWithKey' #-}
1481
1482 -- | /O(n)/. Fold the keys and values in the map using the given left-associative
1483 -- binary operator, such that
1484 -- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@.
1485 --
1486 -- For example,
1487 --
1488 -- > keys = reverse . foldlWithKey (\ks k x -> k:ks) []
1489 --
1490 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
1491 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
1492 foldlWithKey :: (a -> Int -> b -> a) -> a -> IntMap b -> a
1493 foldlWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1494 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1495 | otherwise -> go (go z l) r
1496 _ -> go z t
1497 where
1498 go z' Nil = z'
1499 go z' (Tip kx x) = f z' kx x
1500 go z' (Bin _ _ l r) = go (go z' l) r
1501 {-# INLINE foldlWithKey #-}
1502
1503 -- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is
1504 -- evaluated before using the result in the next application. This
1505 -- function is strict in the starting value.
1506 foldlWithKey' :: (a -> Int -> b -> a) -> a -> IntMap b -> a
1507 foldlWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only.
1508 case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before
1509 | otherwise -> go (go z l) r
1510 _ -> go z t
1511 where
1512 STRICT_1_OF_2(go)
1513 go z' Nil = z'
1514 go z' (Tip kx x) = f z' kx x
1515 go z' (Bin _ _ l r) = go (go z' l) r
1516 {-# INLINE foldlWithKey' #-}
1517
1518 {--------------------------------------------------------------------
1519 List variations
1520 --------------------------------------------------------------------}
1521 -- | /O(n)/.
1522 -- Return all elements of the map in the ascending order of their keys.
1523 -- Subject to list fusion.
1524 --
1525 -- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
1526 -- > elems empty == []
1527
1528 elems :: IntMap a -> [a]
1529 elems = foldr (:) []
1530
1531 -- | /O(n)/. Return all keys of the map in ascending order. Subject to list
1532 -- fusion.
1533 --
1534 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
1535 -- > keys empty == []
1536
1537 keys :: IntMap a -> [Key]
1538 keys = foldrWithKey (\k _ ks -> k : ks) []
1539
1540 -- | /O(n*min(n,W))/. The set of all keys of the map.
1541 --
1542 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
1543 -- > keysSet empty == Data.IntSet.empty
1544
1545 keysSet :: IntMap a -> IntSet.IntSet
1546 keysSet m = IntSet.fromDistinctAscList (keys m)
1547
1548
1549 -- | /O(n)/. An alias for 'toAscList'. Returns all key\/value pairs in the
1550 -- map in ascending key order. Subject to list fusion.
1551 --
1552 -- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1553 -- > assocs empty == []
1554
1555 assocs :: IntMap a -> [(Key,a)]
1556 assocs = toAscList
1557
1558
1559 {--------------------------------------------------------------------
1560 Lists
1561 --------------------------------------------------------------------}
1562 -- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list
1563 -- fusion.
1564 --
1565 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1566 -- > toList empty == []
1567
1568 toList :: IntMap a -> [(Key,a)]
1569 toList = toAscList
1570
1571 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
1572 -- keys are in ascending order. Subject to list fusion.
1573 --
1574 -- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
1575
1576 toAscList :: IntMap a -> [(Key,a)]
1577 toAscList = foldrWithKey (\k x xs -> (k,x):xs) []
1578
1579 -- | /O(n)/. Convert the map to a list of key\/value pairs where the keys
1580 -- are in descending order. Subject to list fusion.
1581 --
1582 -- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")]
1583
1584 toDescList :: IntMap a -> [(Key,a)]
1585 toDescList = foldlWithKey (\xs k x -> (k,x):xs) []
1586
1587 #if __GLASGOW_HASKELL__
1588 -- List fusion for the list generating functions
1589 {-# RULES "IntMap/elems" forall im . elems im = build (\c n -> foldr c n im) #-}
1590 {-# RULES "IntMap/keys" forall im . keys im = build (\c n -> foldrWithKey (\k _ ks -> c k ks) n im) #-}
1591 {-# RULES "IntMap/assocs" forall im . assocs im = build (\c n -> foldrWithKey (\k x xs -> c (k,x) xs) n im) #-}
1592 {-# RULES "IntMap/toList" forall im . toList im = build (\c n -> foldrWithKey (\k x xs -> c (k,x) xs) n im) #-}
1593 {-# RULES "IntMap/toAscList" forall im . toAscList im = build (\c n -> foldrWithKey (\k x xs -> c (k,x) xs) n im) #-}
1594 {-# RULES "IntMap/toDescList" forall im . toDescList im = build (\c n -> foldlWithKey (\xs k x -> c (k,x) xs) n im) #-}
1595 #endif
1596
1597
1598 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
1599 --
1600 -- > fromList [] == empty
1601 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
1602 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
1603
1604 fromList :: [(Key,a)] -> IntMap a
1605 fromList xs
1606 = foldlStrict ins empty xs
1607 where
1608 ins t (k,x) = insert k x t
1609
1610 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1611 --
1612 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "ab"), (5, "cba")]
1613 -- > fromListWith (++) [] == empty
1614
1615 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1616 fromListWith f xs
1617 = fromListWithKey (\_ x y -> f x y) xs
1618
1619 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
1620 --
1621 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
1622 -- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")]
1623 -- > fromListWithKey f [] == empty
1624
1625 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1626 fromListWithKey f xs
1627 = foldlStrict ins empty xs
1628 where
1629 ins t (k,x) = insertWithKey f k x t
1630
1631 -- | /O(n)/. Build a map from a list of key\/value pairs where
1632 -- the keys are in ascending order.
1633 --
1634 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1635 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
1636
1637 fromAscList :: [(Key,a)] -> IntMap a
1638 fromAscList xs
1639 = fromAscListWithKey (\_ x _ -> x) xs
1640
1641 -- | /O(n)/. Build a map from a list of key\/value pairs where
1642 -- the keys are in ascending order, with a combining function on equal keys.
1643 -- /The precondition (input list is ascending) is not checked./
1644 --
1645 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1646
1647 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1648 fromAscListWith f xs
1649 = fromAscListWithKey (\_ x y -> f x y) xs
1650
1651 -- | /O(n)/. Build a map from a list of key\/value pairs where
1652 -- the keys are in ascending order, with a combining function on equal keys.
1653 -- /The precondition (input list is ascending) is not checked./
1654 --
1655 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
1656 -- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")]
1657
1658 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1659 fromAscListWithKey _ [] = Nil
1660 fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
1661 where
1662 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1663 combineEq z [] = [z]
1664 combineEq z@(kz,zz) (x@(kx,xx):xs)
1665 | kx==kz = let yy = f kx xx zz in combineEq (kx,yy) xs
1666 | otherwise = z:combineEq x xs
1667
1668 -- | /O(n)/. Build a map from a list of key\/value pairs where
1669 -- the keys are in ascending order and all distinct.
1670 -- /The precondition (input list is strictly ascending) is not checked./
1671 --
1672 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1673
1674 fromDistinctAscList :: [(Key,a)] -> IntMap a
1675 fromDistinctAscList [] = Nil
1676 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
1677 where
1678 work (kx,vx) [] stk = finish kx (Tip kx vx) stk
1679 work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
1680
1681 reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
1682 reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
1683 reduce z zs m px tx stk@(Push py ty stk') =
1684 let mxy = branchMask px py
1685 pxy = mask px mxy
1686 in if shorter m mxy
1687 then reduce z zs m pxy (Bin pxy mxy ty tx) stk'
1688 else work z zs (Push px tx stk)
1689
1690 finish _ t Nada = t
1691 finish px tx (Push py ty stk) = finish p (join py ty px tx) stk
1692 where m = branchMask px py
1693 p = mask px m
1694
1695 data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada
1696
1697
1698 {--------------------------------------------------------------------
1699 Eq
1700 --------------------------------------------------------------------}
1701 instance Eq a => Eq (IntMap a) where
1702 t1 == t2 = equal t1 t2
1703 t1 /= t2 = nequal t1 t2
1704
1705 equal :: Eq a => IntMap a -> IntMap a -> Bool
1706 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1707 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
1708 equal (Tip kx x) (Tip ky y)
1709 = (kx == ky) && (x==y)
1710 equal Nil Nil = True
1711 equal _ _ = False
1712
1713 nequal :: Eq a => IntMap a -> IntMap a -> Bool
1714 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1715 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
1716 nequal (Tip kx x) (Tip ky y)
1717 = (kx /= ky) || (x/=y)
1718 nequal Nil Nil = False
1719 nequal _ _ = True
1720
1721 {--------------------------------------------------------------------
1722 Ord
1723 --------------------------------------------------------------------}
1724
1725 instance Ord a => Ord (IntMap a) where
1726 compare m1 m2 = compare (toList m1) (toList m2)
1727
1728 {--------------------------------------------------------------------
1729 Functor
1730 --------------------------------------------------------------------}
1731
1732 instance Functor IntMap where
1733 fmap = map
1734
1735 {--------------------------------------------------------------------
1736 Show
1737 --------------------------------------------------------------------}
1738
1739 instance Show a => Show (IntMap a) where
1740 showsPrec d m = showParen (d > 10) $
1741 showString "fromList " . shows (toList m)
1742
1743 {--------------------------------------------------------------------
1744 Read
1745 --------------------------------------------------------------------}
1746 instance (Read e) => Read (IntMap e) where
1747 #ifdef __GLASGOW_HASKELL__
1748 readPrec = parens $ prec 10 $ do
1749 Ident "fromList" <- lexP
1750 xs <- readPrec
1751 return (fromList xs)
1752
1753 readListPrec = readListPrecDefault
1754 #else
1755 readsPrec p = readParen (p > 10) $ \ r -> do
1756 ("fromList",s) <- lex r
1757 (xs,t) <- reads s
1758 return (fromList xs,t)
1759 #endif
1760
1761 {--------------------------------------------------------------------
1762 Typeable
1763 --------------------------------------------------------------------}
1764
1765 #include "Typeable.h"
1766 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1767
1768 {--------------------------------------------------------------------
1769 Helpers
1770 --------------------------------------------------------------------}
1771 {--------------------------------------------------------------------
1772 Join
1773 --------------------------------------------------------------------}
1774 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1775 join p1 t1 p2 t2
1776 | zero p1 m = Bin p m t1 t2
1777 | otherwise = Bin p m t2 t1
1778 where
1779 m = branchMask p1 p2
1780 p = mask p1 m
1781 {-# INLINE join #-}
1782
1783 {--------------------------------------------------------------------
1784 @bin@ assures that we never have empty trees within a tree.
1785 --------------------------------------------------------------------}
1786 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1787 bin _ _ l Nil = l
1788 bin _ _ Nil r = r
1789 bin p m l r = Bin p m l r
1790 {-# INLINE bin #-}
1791
1792
1793 {--------------------------------------------------------------------
1794 Endian independent bit twiddling
1795 --------------------------------------------------------------------}
1796 zero :: Key -> Mask -> Bool
1797 zero i m
1798 = (natFromInt i) .&. (natFromInt m) == 0
1799 {-# INLINE zero #-}
1800
1801 nomatch,match :: Key -> Prefix -> Mask -> Bool
1802 nomatch i p m
1803 = (mask i m) /= p
1804 {-# INLINE nomatch #-}
1805
1806 match i p m
1807 = (mask i m) == p
1808 {-# INLINE match #-}
1809
1810 mask :: Key -> Mask -> Prefix
1811 mask i m
1812 = maskW (natFromInt i) (natFromInt m)
1813 {-# INLINE mask #-}
1814
1815
1816 {--------------------------------------------------------------------
1817 Big endian operations
1818 --------------------------------------------------------------------}
1819 maskW :: Nat -> Nat -> Prefix
1820 maskW i m
1821 = intFromNat (i .&. (complement (m-1) `xor` m))
1822 {-# INLINE maskW #-}
1823
1824 shorter :: Mask -> Mask -> Bool
1825 shorter m1 m2
1826 = (natFromInt m1) > (natFromInt m2)
1827 {-# INLINE shorter #-}
1828
1829 branchMask :: Prefix -> Prefix -> Mask
1830 branchMask p1 p2
1831 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1832 {-# INLINE branchMask #-}
1833
1834 {----------------------------------------------------------------------
1835 Finding the highest bit (mask) in a word [x] can be done efficiently in
1836 three ways:
1837 * convert to a floating point value and the mantissa tells us the
1838 [log2(x)] that corresponds with the highest bit position. The mantissa
1839 is retrieved either via the standard C function [frexp] or by some bit
1840 twiddling on IEEE compatible numbers (float). Note that one needs to
1841 use at least [double] precision for an accurate mantissa of 32 bit
1842 numbers.
1843 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1844 * use processor specific assembler instruction (asm).
1845
1846 The most portable way would be [bit], but is it efficient enough?
1847 I have measured the cycle counts of the different methods on an AMD
1848 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1849
1850 highestBitMask: method cycles
1851 --------------
1852 frexp 200
1853 float 33
1854 bit 11
1855 asm 12
1856
1857 highestBit: method cycles
1858 --------------
1859 frexp 195
1860 float 33
1861 bit 11
1862 asm 11
1863
1864 Wow, the bit twiddling is on today's RISC like machines even faster
1865 than a single CISC instruction (BSR)!
1866 ----------------------------------------------------------------------}
1867
1868 {----------------------------------------------------------------------
1869 [highestBitMask] returns a word where only the highest bit is set.
1870 It is found by first setting all bits in lower positions than the
1871 highest bit and than taking an exclusive or with the original value.
1872 Allthough the function may look expensive, GHC compiles this into
1873 excellent C code that subsequently compiled into highly efficient
1874 machine code. The algorithm is derived from Jorg Arndt's FXT library.
1875 ----------------------------------------------------------------------}
1876 highestBitMask :: Nat -> Nat
1877 highestBitMask x0
1878 = case (x0 .|. shiftRL x0 1) of
1879 x1 -> case (x1 .|. shiftRL x1 2) of
1880 x2 -> case (x2 .|. shiftRL x2 4) of
1881 x3 -> case (x3 .|. shiftRL x3 8) of
1882 x4 -> case (x4 .|. shiftRL x4 16) of
1883 x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms
1884 x6 -> (x6 `xor` (shiftRL x6 1))
1885 {-# INLINE highestBitMask #-}
1886
1887
1888 {--------------------------------------------------------------------
1889 Utilities
1890 --------------------------------------------------------------------}
1891
1892 foldlStrict :: (a -> b -> a) -> a -> [b] -> a
1893 foldlStrict f = go
1894 where
1895 go z [] = z
1896 go z (x:xs) = let z' = f z x in z' `seq` go z' xs
1897 {-# INLINE foldlStrict #-}
1898
1899 {--------------------------------------------------------------------
1900 Debugging
1901 --------------------------------------------------------------------}
1902 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1903 -- in a compressed, hanging format.
1904 showTree :: Show a => IntMap a -> String
1905 showTree s
1906 = showTreeWith True False s
1907
1908
1909 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1910 the tree that implements the map. If @hang@ is
1911 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1912 @wide@ is 'True', an extra wide version is shown.
1913 -}
1914 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1915 showTreeWith hang wide t
1916 | hang = (showsTreeHang wide [] t) ""
1917 | otherwise = (showsTree wide [] [] t) ""
1918
1919 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1920 showsTree wide lbars rbars t
1921 = case t of
1922 Bin p m l r
1923 -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1924 showWide wide rbars .
1925 showsBars lbars . showString (showBin p m) . showString "\n" .
1926 showWide wide lbars .
1927 showsTree wide (withEmpty lbars) (withBar lbars) l
1928 Tip k x
1929 -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1930 Nil -> showsBars lbars . showString "|\n"
1931
1932 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1933 showsTreeHang wide bars t
1934 = case t of
1935 Bin p m l r
1936 -> showsBars bars . showString (showBin p m) . showString "\n" .
1937 showWide wide bars .
1938 showsTreeHang wide (withBar bars) l .
1939 showWide wide bars .
1940 showsTreeHang wide (withEmpty bars) r
1941 Tip k x
1942 -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1943 Nil -> showsBars bars . showString "|\n"
1944
1945 showBin :: Prefix -> Mask -> String
1946 showBin _ _
1947 = "*" -- ++ show (p,m)
1948
1949 showWide :: Bool -> [String] -> String -> String
1950 showWide wide bars
1951 | wide = showString (concat (reverse bars)) . showString "|\n"
1952 | otherwise = id
1953
1954 showsBars :: [String] -> ShowS
1955 showsBars bars
1956 = case bars of
1957 [] -> id
1958 _ -> showString (concat (reverse (tail bars))) . showString node
1959
1960 node :: String
1961 node = "+--"
1962
1963 withBar, withEmpty :: [String] -> [String]
1964 withBar bars = "| ":bars
1965 withEmpty bars = " ":bars