Make strict IntMap merges strict
[packages/containers.git] / Data / IntMap / Strict / Internal.hs
1 {-# LANGUAGE CPP #-}
2 {-# LANGUAGE BangPatterns #-}
3
4 #include "containers.h"
5
6 -----------------------------------------------------------------------------
7 -- |
8 -- Module : Data.IntMap.Strict.Internal
9 -- Copyright : (c) Daan Leijen 2002
10 -- (c) Andriy Palamarchuk 2008
11 -- License : BSD-style
12 -- Maintainer : libraries@haskell.org
13 -- Portability : portable
14 --
15 --
16 -- = Finite Int Maps (strict interface)
17 --
18 -- The @'IntMap' v@ type represents a finite map (sometimes called a dictionary)
19 -- from key of type @Int@ to values of type @v@.
20 --
21 -- Each function in this module is careful to force values before installing
22 -- them in an 'IntMap'. This is usually more efficient when laziness is not
23 -- necessary. When laziness /is/ required, use the functions in
24 -- "Data.IntMap.Lazy".
25 --
26 -- In particular, the functions in this module obey the following law:
27 --
28 -- - If all values stored in all maps in the arguments are in WHNF, then all
29 -- values stored in all maps in the results will be in WHNF once those maps
30 -- are evaluated.
31 --
32 -- For a walkthrough of the most commonly used functions see the
33 -- <https://haskell-containers.readthedocs.io/en/latest/map.html maps introduction>.
34 --
35 -- This module is intended to be imported qualified, to avoid name clashes with
36 -- Prelude functions:
37 --
38 -- > import Data.IntMap.Strict (IntMap)
39 -- > import qualified Data.IntMap.Strict as IntMap
40 --
41 -- Note that the implementation is generally /left-biased/. Functions that take
42 -- two maps as arguments and combine them, such as `union` and `intersection`,
43 -- prefer the values in the first argument to those in the second.
44 --
45 --
46 -- == Detailed performance information
47 --
48 -- The amortized running time is given for each operation, with /n/ referring to
49 -- the number of entries in the map and /W/ referring to the number of bits in
50 -- an 'Int' (32 or 64).
51 --
52 -- Benchmarks comparing "Data.IntMap.Strict" with other dictionary
53 -- implementations can be found at https://github.com/haskell-perf/dictionaries.
54 --
55 --
56 -- == Warning
57 --
58 -- The 'IntMap' type is shared between the lazy and strict modules, meaning that
59 -- the same 'IntMap' value can be passed to functions in both modules. This
60 -- means that the 'Functor', 'Traversable' and 'Data.Data.Data' instances are
61 -- the same as for the "Data.IntMap.Lazy" module, so if they are used the
62 -- resulting map may contain suspended values (thunks).
63 --
64 --
65 -- == Implementation
66 --
67 -- The implementation is based on /big-endian patricia trees/. This data
68 -- structure performs especially well on binary operations like 'union' and
69 -- 'intersection'. Additionally, benchmarks show that it is also (much) faster
70 -- on insertions and deletions when compared to a generic size-balanced map
71 -- implementation (see "Data.Map").
72 --
73 -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
74 -- Workshop on ML, September 1998, pages 77-86,
75 -- <http://citeseer.ist.psu.edu/okasaki98fast.html>
76 --
77 -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
78 -- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
79 -- October 1968, pages 514-534.
80 --
81 -----------------------------------------------------------------------------
82
83 -- See the notes at the beginning of Data.IntMap.Internal.
84
85 module Data.IntMap.Strict.Internal (
86 -- * Map type
87 #if !defined(TESTING)
88 IntMap, Key -- instance Eq,Show
89 #else
90 IntMap(..), Key -- instance Eq,Show
91 #endif
92
93 -- * Construction
94 , empty
95 , singleton
96 , fromSet
97
98 -- ** From Unordered Lists
99 , fromList
100 , fromListWith
101 , fromListWithKey
102
103 -- ** From Ascending Lists
104 , fromAscList
105 , fromAscListWith
106 , fromAscListWithKey
107 , fromDistinctAscList
108
109 -- * Insertion
110 , insert
111 , insertWith
112 , insertWithKey
113 , insertLookupWithKey
114
115 -- * Deletion\/Update
116 , delete
117 , adjust
118 , adjustWithKey
119 , update
120 , updateWithKey
121 , updateLookupWithKey
122 , alter
123 , alterF
124
125 -- * Query
126 -- ** Lookup
127 , lookup
128 , (!?)
129 , (!)
130 , findWithDefault
131 , member
132 , notMember
133 , lookupLT
134 , lookupGT
135 , lookupLE
136 , lookupGE
137
138 -- ** Size
139 , null
140 , size
141
142 -- * Combine
143
144 -- ** Union
145 , union
146 , unionWith
147 , unionWithKey
148 , unions
149 , unionsWith
150
151 -- ** Difference
152 , difference
153 , (\\)
154 , differenceWith
155 , differenceWithKey
156
157 -- ** Intersection
158 , intersection
159 , intersectionWith
160 , intersectionWithKey
161
162 -- ** Universal combining function
163 , mergeWithKey
164
165 -- * Traversal
166 -- ** Map
167 , map
168 , mapWithKey
169 , traverseWithKey
170 , traverseMaybeWithKey
171 , mapAccum
172 , mapAccumWithKey
173 , mapAccumRWithKey
174 , mapKeys
175 , mapKeysWith
176 , mapKeysMonotonic
177
178 -- * Folds
179 , foldr
180 , foldl
181 , foldrWithKey
182 , foldlWithKey
183 , foldMapWithKey
184
185 -- ** Strict folds
186 , foldr'
187 , foldl'
188 , foldrWithKey'
189 , foldlWithKey'
190
191 -- * Conversion
192 , elems
193 , keys
194 , assocs
195 , keysSet
196
197 -- ** Lists
198 , toList
199
200 -- ** Ordered lists
201 , toAscList
202 , toDescList
203
204 -- * Filter
205 , filter
206 , filterWithKey
207 , restrictKeys
208 , withoutKeys
209 , partition
210 , partitionWithKey
211
212 , mapMaybe
213 , mapMaybeWithKey
214 , mapEither
215 , mapEitherWithKey
216
217 , split
218 , splitLookup
219 , splitRoot
220
221 -- * Submap
222 , isSubmapOf, isSubmapOfBy
223 , isProperSubmapOf, isProperSubmapOfBy
224
225 -- * Min\/Max
226 , lookupMin
227 , lookupMax
228 , findMin
229 , findMax
230 , deleteMin
231 , deleteMax
232 , deleteFindMin
233 , deleteFindMax
234 , updateMin
235 , updateMax
236 , updateMinWithKey
237 , updateMaxWithKey
238 , minView
239 , maxView
240 , minViewWithKey
241 , maxViewWithKey
242
243 #ifdef __GLASGOW_HASKELL__
244 -- * Debugging
245 , showTree
246 , showTreeWith
247 #endif
248 ) where
249
250 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
251
252 import Data.Bits
253 import qualified Data.IntMap.Internal as L
254 import Data.IntMap.Internal
255 ( IntMap (..)
256 , Key
257 , Prefix
258 , Mask
259 , mask
260 , branchMask
261 , shorter
262 , nomatch
263 , zero
264 , natFromInt
265 , intFromNat
266 , bin
267 , binCheckLeft
268 , binCheckRight
269 , link
270
271 , (\\)
272 , (!)
273 , (!?)
274 , empty
275 , assocs
276 , filter
277 , filterWithKey
278 , findMin
279 , findMax
280 , foldMapWithKey
281 , foldr
282 , foldl
283 , foldr'
284 , foldl'
285 , foldlWithKey
286 , foldrWithKey
287 , foldlWithKey'
288 , foldrWithKey'
289 , keysSet
290 , mergeWithKey'
291 , delete
292 , deleteMin
293 , deleteMax
294 , deleteFindMax
295 , deleteFindMin
296 , difference
297 , elems
298 , intersection
299 , isProperSubmapOf
300 , isProperSubmapOfBy
301 , isSubmapOf
302 , isSubmapOfBy
303 , lookup
304 , lookupLE
305 , lookupGE
306 , lookupLT
307 , lookupGT
308 , lookupMin
309 , lookupMax
310 , minView
311 , maxView
312 , minViewWithKey
313 , maxViewWithKey
314 , keys
315 , mapKeys
316 , mapKeysMonotonic
317 , member
318 , notMember
319 , null
320 , partition
321 , partitionWithKey
322 , restrictKeys
323 , size
324 , split
325 , splitLookup
326 , splitRoot
327 , toAscList
328 , toDescList
329 , toList
330 , union
331 , unions
332 , withoutKeys
333 )
334 #ifdef __GLASGOW_HASKELL__
335 import Data.IntMap.Internal.DeprecatedDebug (showTree, showTreeWith)
336 #endif
337 import qualified Data.IntSet.Internal as IntSet
338 import Utils.Containers.Internal.BitUtil
339 import Utils.Containers.Internal.StrictPair
340 #if !MIN_VERSION_base(4,8,0)
341 import Data.Functor((<$>))
342 #endif
343 import Control.Applicative (Applicative (..), liftA2)
344 import qualified Data.Foldable as Foldable
345 #if !MIN_VERSION_base(4,8,0)
346 import Data.Foldable (Foldable())
347 #endif
348
349 {--------------------------------------------------------------------
350 Query
351 --------------------------------------------------------------------}
352
353 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
354 -- returns the value at key @k@ or returns @def@ when the key is not an
355 -- element of the map.
356 --
357 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
358 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
359
360 -- See IntMap.Internal.Note: Local 'go' functions and capturing]
361 findWithDefault :: a -> Key -> IntMap a -> a
362 findWithDefault def !k = go
363 where
364 go (Bin p m l r) | nomatch k p m = def
365 | zero k m = go l
366 | otherwise = go r
367 go (Tip kx x) | k == kx = x
368 | otherwise = def
369 go Nil = def
370
371 {--------------------------------------------------------------------
372 Construction
373 --------------------------------------------------------------------}
374 -- | /O(1)/. A map of one element.
375 --
376 -- > singleton 1 'a' == fromList [(1, 'a')]
377 -- > size (singleton 1 'a') == 1
378
379 singleton :: Key -> a -> IntMap a
380 singleton k !x
381 = Tip k x
382 {-# INLINE singleton #-}
383
384 {--------------------------------------------------------------------
385 Insert
386 --------------------------------------------------------------------}
387 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
388 -- If the key is already present in the map, the associated value is
389 -- replaced with the supplied value, i.e. 'insert' is equivalent to
390 -- @'insertWith' 'const'@.
391 --
392 -- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
393 -- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
394 -- > insert 5 'x' empty == singleton 5 'x'
395
396 insert :: Key -> a -> IntMap a -> IntMap a
397 insert !k !x t =
398 case t of
399 Bin p m l r
400 | nomatch k p m -> link k (Tip k x) p t
401 | zero k m -> Bin p m (insert k x l) r
402 | otherwise -> Bin p m l (insert k x r)
403 Tip ky _
404 | k==ky -> Tip k x
405 | otherwise -> link k (Tip k x) ky t
406 Nil -> Tip k x
407
408 -- right-biased insertion, used by 'union'
409 -- | /O(min(n,W))/. Insert with a combining function.
410 -- @'insertWith' f key value mp@
411 -- will insert the pair (key, value) into @mp@ if key does
412 -- not exist in the map. If the key does exist, the function will
413 -- insert @f new_value old_value@.
414 --
415 -- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
416 -- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
417 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx"
418
419 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
420 insertWith f k x t
421 = insertWithKey (\_ x' y' -> f x' y') k x t
422
423 -- | /O(min(n,W))/. Insert with a combining function.
424 -- @'insertWithKey' f key value mp@
425 -- will insert the pair (key, value) into @mp@ if key does
426 -- not exist in the map. If the key does exist, the function will
427 -- insert @f key new_value old_value@.
428 --
429 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
430 -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
431 -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
432 -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx"
433 --
434 -- If the key exists in the map, this function is lazy in @value@ but strict
435 -- in the result of @f@.
436
437 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
438 insertWithKey f !k x t =
439 case t of
440 Bin p m l r
441 | nomatch k p m -> link k (singleton k x) p t
442 | zero k m -> Bin p m (insertWithKey f k x l) r
443 | otherwise -> Bin p m l (insertWithKey f k x r)
444 Tip ky y
445 | k==ky -> Tip k $! f k x y
446 | otherwise -> link k (singleton k x) ky t
447 Nil -> singleton k x
448
449 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
450 -- is a pair where the first element is equal to (@'lookup' k map@)
451 -- and the second element equal to (@'insertWithKey' f k x map@).
452 --
453 -- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
454 -- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
455 -- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")])
456 -- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx")
457 --
458 -- This is how to define @insertLookup@ using @insertLookupWithKey@:
459 --
460 -- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t
461 -- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")])
462 -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")])
463
464 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
465 insertLookupWithKey f0 !k0 x0 t0 = toPair $ go f0 k0 x0 t0
466 where
467 go f k x t =
468 case t of
469 Bin p m l r
470 | nomatch k p m -> Nothing :*: link k (singleton k x) p t
471 | zero k m -> let (found :*: l') = go f k x l in (found :*: Bin p m l' r)
472 | otherwise -> let (found :*: r') = go f k x r in (found :*: Bin p m l r')
473 Tip ky y
474 | k==ky -> (Just y :*: (Tip k $! f k x y))
475 | otherwise -> (Nothing :*: link k (singleton k x) ky t)
476 Nil -> Nothing :*: (singleton k x)
477
478
479 {--------------------------------------------------------------------
480 Deletion
481 --------------------------------------------------------------------}
482 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
483 -- a member of the map, the original map is returned.
484 --
485 -- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
486 -- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
487 -- > adjust ("new " ++) 7 empty == empty
488
489 adjust :: (a -> a) -> Key -> IntMap a -> IntMap a
490 adjust f k m
491 = adjustWithKey (\_ x -> f x) k m
492
493 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
494 -- a member of the map, the original map is returned.
495 --
496 -- > let f key x = (show key) ++ ":new " ++ x
497 -- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
498 -- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
499 -- > adjustWithKey f 7 empty == empty
500
501 adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a
502 adjustWithKey f !k t =
503 case t of
504 Bin p m l r
505 | nomatch k p m -> t
506 | zero k m -> Bin p m (adjustWithKey f k l) r
507 | otherwise -> Bin p m l (adjustWithKey f k r)
508 Tip ky y
509 | k==ky -> Tip ky $! f k y
510 | otherwise -> t
511 Nil -> Nil
512
513 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
514 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
515 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
516 --
517 -- > let f x = if x == "a" then Just "new a" else Nothing
518 -- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
519 -- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
520 -- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
521
522 update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a
523 update f
524 = updateWithKey (\_ x -> f x)
525
526 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
527 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
528 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
529 --
530 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
531 -- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
532 -- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
533 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
534
535 updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
536 updateWithKey f !k t =
537 case t of
538 Bin p m l r
539 | nomatch k p m -> t
540 | zero k m -> binCheckLeft p m (updateWithKey f k l) r
541 | otherwise -> binCheckRight p m l (updateWithKey f k r)
542 Tip ky y
543 | k==ky -> case f k y of
544 Just !y' -> Tip ky y'
545 Nothing -> Nil
546 | otherwise -> t
547 Nil -> Nil
548
549 -- | /O(min(n,W))/. Lookup and update.
550 -- The function returns original value, if it is updated.
551 -- This is different behavior than 'Data.Map.updateLookupWithKey'.
552 -- Returns the original key value if the map entry is deleted.
553 --
554 -- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
555 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
556 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")])
557 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
558
559 updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
560 updateLookupWithKey f0 !k0 t0 = toPair $ go f0 k0 t0
561 where
562 go f k t =
563 case t of
564 Bin p m l r
565 | nomatch k p m -> (Nothing :*: t)
566 | zero k m -> let (found :*: l') = go f k l in (found :*: binCheckLeft p m l' r)
567 | otherwise -> let (found :*: r') = go f k r in (found :*: binCheckRight p m l r')
568 Tip ky y
569 | k==ky -> case f k y of
570 Just !y' -> (Just y :*: Tip ky y')
571 Nothing -> (Just y :*: Nil)
572 | otherwise -> (Nothing :*: t)
573 Nil -> (Nothing :*: Nil)
574
575
576
577 -- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
578 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
579 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
580 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
581 alter f !k t =
582 case t of
583 Bin p m l r
584 | nomatch k p m -> case f Nothing of
585 Nothing -> t
586 Just !x -> link k (Tip k x) p t
587 | zero k m -> binCheckLeft p m (alter f k l) r
588 | otherwise -> binCheckRight p m l (alter f k r)
589 Tip ky y
590 | k==ky -> case f (Just y) of
591 Just !x -> Tip ky x
592 Nothing -> Nil
593 | otherwise -> case f Nothing of
594 Just !x -> link k (Tip k x) ky t
595 Nothing -> t
596 Nil -> case f Nothing of
597 Just !x -> Tip k x
598 Nothing -> Nil
599
600 -- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at
601 -- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete,
602 -- or update a value in an 'IntMap'. In short : @'lookup' k <$> 'alterF' f k m = f
603 -- ('lookup' k m)@.
604 --
605 -- Example:
606 --
607 -- @
608 -- interactiveAlter :: Int -> IntMap String -> IO (IntMap String)
609 -- interactiveAlter k m = alterF f k m where
610 -- f Nothing = do
611 -- putStrLn $ show k ++
612 -- " was not found in the map. Would you like to add it?"
613 -- getUserResponse1 :: IO (Maybe String)
614 -- f (Just old) = do
615 -- putStrLn $ "The key is currently bound to " ++ show old ++
616 -- ". Would you like to change or delete it?"
617 -- getUserResponse2 :: IO (Maybe String)
618 -- @
619 --
620 -- 'alterF' is the most general operation for working with an individual
621 -- key that may or may not be in a given map.
622
623 -- Note: 'alterF' is a flipped version of the 'at' combinator from
624 -- 'Control.Lens.At'.
625 --
626 -- @since 0.5.8
627
628 alterF :: Functor f
629 => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a)
630 -- This implementation was modified from 'Control.Lens.At'.
631 alterF f k m = (<$> f mv) $ \fres ->
632 case fres of
633 Nothing -> maybe m (const (delete k m)) mv
634 Just !v' -> insert k v' m
635 where mv = lookup k m
636
637
638 {--------------------------------------------------------------------
639 Union
640 --------------------------------------------------------------------}
641 -- | The union of a list of maps, with a combining operation.
642 --
643 -- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])]
644 -- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")]
645
646 unionsWith :: Foldable f => (a->a->a) -> f (IntMap a) -> IntMap a
647 unionsWith f ts
648 = Foldable.foldl' (unionWith f) empty ts
649
650 -- | /O(n+m)/. The union with a combining function.
651 --
652 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
653
654 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
655 unionWith f m1 m2
656 = unionWithKey (\_ x y -> f x y) m1 m2
657
658 -- | /O(n+m)/. The union with a combining function.
659 --
660 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
661 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
662
663 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
664 unionWithKey f m1 m2
665 = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 $! f k1 x1 x2) id id m1 m2
666
667 {--------------------------------------------------------------------
668 Difference
669 --------------------------------------------------------------------}
670
671 -- | /O(n+m)/. Difference with a combining function.
672 --
673 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
674 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
675 -- > == singleton 3 "b:B"
676
677 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
678 differenceWith f m1 m2
679 = differenceWithKey (\_ x y -> f x y) m1 m2
680
681 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
682 -- encountered, the combining function is applied to the key and both values.
683 -- If it returns 'Nothing', the element is discarded (proper set difference).
684 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
685 --
686 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
687 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
688 -- > == singleton 3 "3:b|B"
689
690 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
691 differenceWithKey f m1 m2
692 = mergeWithKey f id (const Nil) m1 m2
693
694 {--------------------------------------------------------------------
695 Intersection
696 --------------------------------------------------------------------}
697
698 -- | /O(n+m)/. The intersection with a combining function.
699 --
700 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
701
702 intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
703 intersectionWith f m1 m2
704 = intersectionWithKey (\_ x y -> f x y) m1 m2
705
706 -- | /O(n+m)/. The intersection with a combining function.
707 --
708 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
709 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
710
711 intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
712 intersectionWithKey f m1 m2
713 = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 $! f k1 x1 x2) (const Nil) (const Nil) m1 m2
714
715 {--------------------------------------------------------------------
716 MergeWithKey
717 --------------------------------------------------------------------}
718
719 -- | /O(n+m)/. A high-performance universal combining function. Using
720 -- 'mergeWithKey', all combining functions can be defined without any loss of
721 -- efficiency (with exception of 'union', 'difference' and 'intersection',
722 -- where sharing of some nodes is lost with 'mergeWithKey').
723 --
724 -- Please make sure you know what is going on when using 'mergeWithKey',
725 -- otherwise you can be surprised by unexpected code growth or even
726 -- corruption of the data structure.
727 --
728 -- When 'mergeWithKey' is given three arguments, it is inlined to the call
729 -- site. You should therefore use 'mergeWithKey' only to define your custom
730 -- combining functions. For example, you could define 'unionWithKey',
731 -- 'differenceWithKey' and 'intersectionWithKey' as
732 --
733 -- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2
734 -- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2
735 -- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2
736 --
737 -- When calling @'mergeWithKey' combine only1 only2@, a function combining two
738 -- 'IntMap's is created, such that
739 --
740 -- * if a key is present in both maps, it is passed with both corresponding
741 -- values to the @combine@ function. Depending on the result, the key is either
742 -- present in the result with specified value, or is left out;
743 --
744 -- * a nonempty subtree present only in the first map is passed to @only1@ and
745 -- the output is added to the result;
746 --
747 -- * a nonempty subtree present only in the second map is passed to @only2@ and
748 -- the output is added to the result.
749 --
750 -- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/.
751 -- The values can be modified arbitrarily. Most common variants of @only1@ and
752 -- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or
753 -- @'filterWithKey' f@ could be used for any @f@.
754
755 mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c)
756 -> IntMap a -> IntMap b -> IntMap c
757 mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2
758 where -- We use the lambda form to avoid non-exhaustive pattern matches warning.
759 combine = \(Tip k1 x1) (Tip _k2 x2) -> case f k1 x1 x2 of Nothing -> Nil
760 Just !x -> Tip k1 x
761 {-# INLINE combine #-}
762 {-# INLINE mergeWithKey #-}
763
764 {--------------------------------------------------------------------
765 Min\/Max
766 --------------------------------------------------------------------}
767
768 -- | /O(log n)/. Update the value at the minimal key.
769 --
770 -- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
771 -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
772
773 updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
774 updateMinWithKey f t =
775 case t of Bin p m l r | m < 0 -> binCheckRight p m l (go f r)
776 _ -> go f t
777 where
778 go f' (Bin p m l r) = binCheckLeft p m (go f' l) r
779 go f' (Tip k y) = case f' k y of
780 Just !y' -> Tip k y'
781 Nothing -> Nil
782 go _ Nil = error "updateMinWithKey Nil"
783
784 -- | /O(log n)/. Update the value at the maximal key.
785 --
786 -- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
787 -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
788
789 updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
790 updateMaxWithKey f t =
791 case t of Bin p m l r | m < 0 -> binCheckLeft p m (go f l) r
792 _ -> go f t
793 where
794 go f' (Bin p m l r) = binCheckRight p m l (go f' r)
795 go f' (Tip k y) = case f' k y of
796 Just !y' -> Tip k y'
797 Nothing -> Nil
798 go _ Nil = error "updateMaxWithKey Nil"
799
800 -- | /O(log n)/. Update the value at the maximal key.
801 --
802 -- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
803 -- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
804
805 updateMax :: (a -> Maybe a) -> IntMap a -> IntMap a
806 updateMax f = updateMaxWithKey (const f)
807
808 -- | /O(log n)/. Update the value at the minimal key.
809 --
810 -- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
811 -- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
812
813 updateMin :: (a -> Maybe a) -> IntMap a -> IntMap a
814 updateMin f = updateMinWithKey (const f)
815
816
817 {--------------------------------------------------------------------
818 Mapping
819 --------------------------------------------------------------------}
820 -- | /O(n)/. Map a function over all values in the map.
821 --
822 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
823
824 map :: (a -> b) -> IntMap a -> IntMap b
825 map f = go
826 where
827 go (Bin p m l r) = Bin p m (go l) (go r)
828 go (Tip k x) = Tip k $! f x
829 go Nil = Nil
830
831 #ifdef __GLASGOW_HASKELL__
832 {-# NOINLINE [1] map #-}
833 {-# RULES
834 "map/map" forall f g xs . map f (map g xs) = map (\x -> f $! g x) xs
835 "map/mapL" forall f g xs . map f (L.map g xs) = map (\x -> f (g x)) xs
836 #-}
837 #endif
838
839 -- | /O(n)/. Map a function over all values in the map.
840 --
841 -- > let f key x = (show key) ++ ":" ++ x
842 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
843
844 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
845 mapWithKey f t
846 = case t of
847 Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
848 Tip k x -> Tip k $! f k x
849 Nil -> Nil
850
851 #ifdef __GLASGOW_HASKELL__
852 -- Pay close attention to strictness here. We need to force the
853 -- intermediate result for map f . map g, and we need to refrain
854 -- from forcing it for map f . L.map g, etc.
855 --
856 -- TODO Consider moving map and mapWithKey to IntMap.Internal so we can write
857 -- non-orphan RULES for things like L.map f (map g xs). We'd need a new function
858 -- for this, and we'd have to pay attention to simplifier phases. Something like
859 --
860 -- lsmap :: (b -> c) -> (a -> b) -> IntMap a -> IntMap c
861 -- lsmap _ _ Nil = Nil
862 -- lsmap f g (Tip k x) = let !gx = g x in Tip k (f gx)
863 -- lsmap f g (Bin p m l r) = Bin p m (lsmap f g l) (lsmap f g r)
864 {-# NOINLINE [1] mapWithKey #-}
865 {-# RULES
866 "mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) =
867 mapWithKey (\k a -> f k $! g k a) xs
868 "mapWithKey/mapWithKeyL" forall f g xs . mapWithKey f (L.mapWithKey g xs) =
869 mapWithKey (\k a -> f k (g k a)) xs
870 "mapWithKey/map" forall f g xs . mapWithKey f (map g xs) =
871 mapWithKey (\k a -> f k $! g a) xs
872 "mapWithKey/mapL" forall f g xs . mapWithKey f (L.map g xs) =
873 mapWithKey (\k a -> f k (g a)) xs
874 "map/mapWithKey" forall f g xs . map f (mapWithKey g xs) =
875 mapWithKey (\k a -> f $! g k a) xs
876 "map/mapWithKeyL" forall f g xs . map f (L.mapWithKey g xs) =
877 mapWithKey (\k a -> f (g k a)) xs
878 #-}
879 #endif
880
881 -- | /O(n)/.
882 -- @'traverseWithKey' f s == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
883 -- That is, behaves exactly like a regular 'traverse' except that the traversing
884 -- function also has access to the key associated with a value.
885 --
886 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')])
887 -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing
888 traverseWithKey :: Applicative t => (Key -> a -> t b) -> IntMap a -> t (IntMap b)
889 traverseWithKey f = go
890 where
891 go Nil = pure Nil
892 go (Tip k v) = (\ !v' -> Tip k v') <$> f k v
893 go (Bin p m l r) = liftA2 (Bin p m) (go l) (go r)
894 {-# INLINE traverseWithKey #-}
895
896 -- | /O(n)/. Traverse keys\/values and collect the 'Just' results.
897 traverseMaybeWithKey
898 :: Applicative f => (Key -> a -> f (Maybe b)) -> IntMap a -> f (IntMap b)
899 traverseMaybeWithKey f = go
900 where
901 go Nil = pure Nil
902 go (Tip k x) = maybe Nil (Tip k $!) <$> f k x
903 go (Bin p m l r) = liftA2 (bin p m) (go l) (go r)
904
905 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
906 -- argument through the map in ascending order of keys.
907 --
908 -- > let f a b = (a ++ b, b ++ "X")
909 -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")])
910
911 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
912 mapAccum f = mapAccumWithKey (\a' _ x -> f a' x)
913
914 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
915 -- argument through the map in ascending order of keys.
916 --
917 -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X")
918 -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")])
919
920 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
921 mapAccumWithKey f a t
922 = mapAccumL f a t
923
924 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
925 -- argument through the map in ascending order of keys. Strict in
926 -- the accumulating argument and the both elements of the
927 -- result of the function.
928 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
929 mapAccumL f0 a0 t0 = toPair $ go f0 a0 t0
930 where
931 go f a t
932 = case t of
933 Bin p m l r -> let (a1 :*: l') = go f a l
934 (a2 :*: r') = go f a1 r
935 in (a2 :*: Bin p m l' r')
936 Tip k x -> let !(a',!x') = f a k x in (a' :*: Tip k x')
937 Nil -> (a :*: Nil)
938
939 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
940 -- argument through the map in descending order of keys.
941 mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
942 mapAccumRWithKey f0 a0 t0 = toPair $ go f0 a0 t0
943 where
944 go f a t
945 = case t of
946 Bin p m l r -> let (a1 :*: r') = go f a r
947 (a2 :*: l') = go f a1 l
948 in (a2 :*: Bin p m l' r')
949 Tip k x -> let !(a',!x') = f a k x in (a' :*: Tip k x')
950 Nil -> (a :*: Nil)
951
952 -- | /O(n*log n)/.
953 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
954 --
955 -- The size of the result may be smaller if @f@ maps two or more distinct
956 -- keys to the same new key. In this case the associated values will be
957 -- combined using @c@.
958 --
959 -- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab"
960 -- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab"
961
962 mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a
963 mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
964
965 {--------------------------------------------------------------------
966 Filter
967 --------------------------------------------------------------------}
968 -- | /O(n)/. Map values and collect the 'Just' results.
969 --
970 -- > let f x = if x == "a" then Just "new a" else Nothing
971 -- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
972
973 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
974 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
975
976 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
977 --
978 -- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
979 -- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
980
981 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
982 mapMaybeWithKey f (Bin p m l r)
983 = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
984 mapMaybeWithKey f (Tip k x) = case f k x of
985 Just !y -> Tip k y
986 Nothing -> Nil
987 mapMaybeWithKey _ Nil = Nil
988
989 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
990 --
991 -- > let f a = if a < "c" then Left a else Right a
992 -- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
993 -- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
994 -- >
995 -- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
996 -- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
997
998 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
999 mapEither f m
1000 = mapEitherWithKey (\_ x -> f x) m
1001
1002 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
1003 --
1004 -- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
1005 -- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1006 -- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
1007 -- >
1008 -- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
1009 -- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
1010
1011 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
1012 mapEitherWithKey f0 t0 = toPair $ go f0 t0
1013 where
1014 go f (Bin p m l r)
1015 = bin p m l1 r1 :*: bin p m l2 r2
1016 where
1017 (l1 :*: l2) = go f l
1018 (r1 :*: r2) = go f r
1019 go f (Tip k x) = case f k x of
1020 Left !y -> (Tip k y :*: Nil)
1021 Right !z -> (Nil :*: Tip k z)
1022 go _ Nil = (Nil :*: Nil)
1023
1024 {--------------------------------------------------------------------
1025 Conversions
1026 --------------------------------------------------------------------}
1027
1028 -- | /O(n)/. Build a map from a set of keys and a function which for each key
1029 -- computes its value.
1030 --
1031 -- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
1032 -- > fromSet undefined Data.IntSet.empty == empty
1033
1034 fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a
1035 fromSet _ IntSet.Nil = Nil
1036 fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r)
1037 fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1)
1038 where -- This is slightly complicated, as we to convert the dense
1039 -- representation of IntSet into tree representation of IntMap.
1040 --
1041 -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'.
1042 -- We split bmask into halves corresponding to left and right subtree.
1043 -- If they are both nonempty, we create a Bin node, otherwise exactly
1044 -- one of them is nonempty and we construct the IntMap from that half.
1045 buildTree g !prefix !bmask bits = case bits of
1046 0 -> Tip prefix $! g prefix
1047 _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of
1048 bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 ->
1049 buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2
1050 | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 ->
1051 buildTree g prefix bmask bits2
1052 | otherwise ->
1053 Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2)
1054
1055 {--------------------------------------------------------------------
1056 Lists
1057 --------------------------------------------------------------------}
1058 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
1059 --
1060 -- > fromList [] == empty
1061 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
1062 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
1063
1064 fromList :: [(Key,a)] -> IntMap a
1065 fromList xs
1066 = Foldable.foldl' ins empty xs
1067 where
1068 ins t (k,x) = insert k x t
1069
1070 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1071 --
1072 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
1073 -- > fromListWith (++) [] == empty
1074
1075 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1076 fromListWith f xs
1077 = fromListWithKey (\_ x y -> f x y) xs
1078
1079 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
1080 --
1081 -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
1082 -- > fromListWith (++) [] == empty
1083
1084 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1085 fromListWithKey f xs
1086 = Foldable.foldl' ins empty xs
1087 where
1088 ins t (k,x) = insertWithKey f k x t
1089
1090 -- | /O(n)/. Build a map from a list of key\/value pairs where
1091 -- the keys are in ascending order.
1092 --
1093 -- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1094 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
1095
1096 fromAscList :: [(Key,a)] -> IntMap a
1097 fromAscList xs
1098 = fromAscListWithKey (\_ x _ -> x) xs
1099
1100 -- | /O(n)/. Build a map from a list of key\/value pairs where
1101 -- the keys are in ascending order, with a combining function on equal keys.
1102 -- /The precondition (input list is ascending) is not checked./
1103 --
1104 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1105
1106 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1107 fromAscListWith f xs
1108 = fromAscListWithKey (\_ x y -> f x y) xs
1109
1110 -- | /O(n)/. Build a map from a list of key\/value pairs where
1111 -- the keys are in ascending order, with a combining function on equal keys.
1112 -- /The precondition (input list is ascending) is not checked./
1113 --
1114 -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
1115
1116 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1117 fromAscListWithKey _ [] = Nil
1118 fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
1119 where
1120 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1121 combineEq z [] = [z]
1122 combineEq z@(kz,zz) (x@(kx,xx):xs)
1123 | kx==kz = let !yy = f kx xx zz in combineEq (kx,yy) xs
1124 | otherwise = z:combineEq x xs
1125
1126 -- | /O(n)/. Build a map from a list of key\/value pairs where
1127 -- the keys are in ascending order and all distinct.
1128 -- /The precondition (input list is strictly ascending) is not checked./
1129 --
1130 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
1131
1132 fromDistinctAscList :: [(Key,a)] -> IntMap a
1133 fromDistinctAscList [] = Nil
1134 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
1135 where
1136 work (kx,!vx) [] stk = finish kx (Tip kx vx) stk
1137 work (kx,!vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
1138
1139 reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
1140 reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
1141 reduce z zs m px tx stk@(Push py ty stk') =
1142 let mxy = branchMask px py
1143 pxy = mask px mxy
1144 in if shorter m mxy
1145 then reduce z zs m pxy (Bin pxy mxy ty tx) stk'
1146 else work z zs (Push px tx stk)
1147
1148 finish _ t Nada = t
1149 finish px tx (Push py ty stk) = finish p (link py ty px tx) stk
1150 where m = branchMask px py
1151 p = mask px m
1152
1153 data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada