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