Data/IntMap.hs

   1 {-# OPTIONS -cpp -fglasgow-exts -fno-bang-patterns #-}
   2 -----------------------------------------------------------------------------
   3 -- |
   4 -- Module      :  Data.IntMap
   5 -- Copyright   :  (c) Daan Leijen 2002
   6 -- License     :  BSD-style
   7 -- Maintainer  :  libraries@haskell.org
   8 -- Stability   :  provisional
   9 -- Portability :  portable
  10 --
  11 -- An efficient implementation of maps from integer keys to values.
  12 --
  13 -- This module is intended to be imported @qualified@, to avoid name
  14 -- clashes with "Prelude" functions.  eg.
  15 --
  16 -- >  import Data.IntMap as Map
  17 --
  18 -- The implementation is based on /big-endian patricia trees/.  This data
  19 -- structure performs especially well on binary operations like 'union'
  20 -- and 'intersection'.  However, my benchmarks show that it is also
  21 -- (much) faster on insertions and deletions when compared to a generic
  22 -- size-balanced map implementation (see "Data.Map" and "Data.FiniteMap").
  23 --
  24 --    * Chris Okasaki and Andy Gill,  \"/Fast Mergeable Integer Maps/\",
  25 --      Workshop on ML, September 1998, pages 77-86,
  26 --      <http://www.cse.ogi.edu/~andy/pub/finite.htm>
  27 --
  28 --    * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
  29 --      Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
  30 --      October 1968, pages 514-534.
  31 --
  32 -- Many operations have a worst-case complexity of /O(min(n,W))/.
  33 -- This means that the operation can become linear in the number of
  34 -- elements with a maximum of /W/ -- the number of bits in an 'Int'
  35 -- (32 or 64).
  36 -----------------------------------------------------------------------------
  37
  38 module Data.IntMap  (
  39             -- * Map type
  40               IntMap, Key          -- instance Eq,Show
  41
  42             -- * Operators
  43             , (!), (\\)
  44
  45             -- * Query
  46             , null
  47             , size
  48             , member
  49             , notMember
  50             , lookup
  51             , findWithDefault
  52
  53             -- * Construction
  54             , empty
  55             , singleton
  56
  57             -- ** Insertion
  58             , insert
  59             , insertWith, insertWithKey, insertLookupWithKey
  60
  61             -- ** Delete\/Update
  62             , delete
  63             , adjust
  64             , adjustWithKey
  65             , update
  66             , updateWithKey
  67             , updateLookupWithKey
  68             , alter
  69
  70             -- * Combine
  71
  72             -- ** Union
  73             , union
  74             , unionWith
  75             , unionWithKey
  76             , unions
  77             , unionsWith
  78
  79             -- ** Difference
  80             , difference
  81             , differenceWith
  82             , differenceWithKey
  83
  84             -- ** Intersection
  85             , intersection
  86             , intersectionWith
  87             , intersectionWithKey
  88
  89             -- * Traversal
  90             -- ** Map
  91             , map
  92             , mapWithKey
  93             , mapAccum
  94             , mapAccumWithKey
  95
  96             -- ** Fold
  97             , fold
  98             , foldWithKey
  99
 100             -- * Conversion
 101             , elems
 102             , keys
 103             , keysSet
 104             , assocs
 105
 106             -- ** Lists
 107             , toList
 108             , fromList
 109             , fromListWith
 110             , fromListWithKey
 111
 112             -- ** Ordered lists
 113             , toAscList
 114             , fromAscList
 115             , fromAscListWith
 116             , fromAscListWithKey
 117             , fromDistinctAscList
 118
 119             -- * Filter
 120             , filter
 121             , filterWithKey
 122             , partition
 123             , partitionWithKey
 124
 125             , mapMaybe
 126             , mapMaybeWithKey
 127             , mapEither
 128             , mapEitherWithKey
 129
 130             , split
 131             , splitLookup
 132
 133             -- * Submap
 134             , isSubmapOf, isSubmapOfBy
 135             , isProperSubmapOf, isProperSubmapOfBy
 136
 137             -- * Debugging
 138             , showTree
 139             , showTreeWith
 140             ) where
 141
 142
 143 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
 144 import Data.Bits
 145 import Data.Int
 146 import qualified Data.IntSet as IntSet
 147 import Data.Monoid (Monoid(..))
 148 import Data.Typeable
 149 import Data.Foldable (Foldable(foldMap))
 150
 151 {-
 152 -- just for testing
 153 import qualified Prelude
 154 import Debug.QuickCheck
 155 import List (nub,sort)
 156 import qualified List
 157 -}
 158
 159 #if __GLASGOW_HASKELL__
 160 import Text.Read
 161 import Data.Generics.Basics
 162 import Data.Generics.Instances
 163 #endif
 164
 165 #if __GLASGOW_HASKELL__ >= 503
 166 import GHC.Word
 167 import GHC.Exts ( Word(..), Int(..), shiftRL# )
 168 #elif __GLASGOW_HASKELL__
 169 import Word
 170 import GlaExts ( Word(..), Int(..), shiftRL# )
 171 #else
 172 import Data.Word
 173 #endif
 174
 175 infixl 9 \\{-This comment teaches CPP correct behaviour -}
 176
 177 -- A "Nat" is a natural machine word (an unsigned Int)
 178 type Nat = Word
 179
 180 natFromInt :: Key -> Nat
 181 natFromInt i = fromIntegral i
 182
 183 intFromNat :: Nat -> Key
 184 intFromNat w = fromIntegral w
 185
 186 shiftRL :: Nat -> Key -> Nat
 187 #if __GLASGOW_HASKELL__
 188 {--------------------------------------------------------------------
 189   GHC: use unboxing to get @shiftRL@ inlined.
 190 --------------------------------------------------------------------}
 191 shiftRL (W# x) (I# i)
 192   = W# (shiftRL# x i)
 193 #else
 194 shiftRL x i   = shiftR x i
 195 #endif
 196
 197 {--------------------------------------------------------------------
 198   Operators
 199 --------------------------------------------------------------------}
 200
 201 -- | /O(min(n,W))/. Find the value at a key.
 202 -- Calls 'error' when the element can not be found.
 203
 204 (!) :: IntMap a -> Key -> a
 205 m ! k    = find' k m
 206
 207 -- | /O(n+m)/. See 'difference'.
 208 (\\) :: IntMap a -> IntMap b -> IntMap a
 209 m1 \\ m2 = difference m1 m2
 210
 211 {--------------------------------------------------------------------
 212   Types
 213 --------------------------------------------------------------------}
 214 -- | A map of integers to values @a@.
 215 data IntMap a = Nil
 216               | Tip {-# UNPACK #-} !Key a
 217               | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a)
 218
 219 type Prefix = Int
 220 type Mask   = Int
 221 type Key    = Int
 222
 223 instance Monoid (IntMap a) where
 224     mempty  = empty
 225     mappend = union
 226     mconcat = unions
 227
 228 instance Foldable IntMap where
 229     foldMap f Nil = mempty
 230     foldMap f (Tip _k v) = f v
 231     foldMap f (Bin _ _ l r) = foldMap f l `mappend` foldMap f r
 232
 233 #if __GLASGOW_HASKELL__
 234
 235 {--------------------------------------------------------------------
 236   A Data instance
 237 --------------------------------------------------------------------}
 238
 239 -- This instance preserves data abstraction at the cost of inefficiency.
 240 -- We omit reflection services for the sake of data abstraction.
 241
 242 instance Data a => Data (IntMap a) where
 243   gfoldl f z im = z fromList `f` (toList im)
 244   toConstr _    = error "toConstr"
 245   gunfold _ _   = error "gunfold"
 246   dataTypeOf _  = mkNorepType "Data.IntMap.IntMap"
 247   dataCast1 f   = gcast1 f
 248
 249 #endif
 250
 251 {--------------------------------------------------------------------
 252   Query
 253 --------------------------------------------------------------------}
 254 -- | /O(1)/. Is the map empty?
 255 null :: IntMap a -> Bool
 256 null Nil   = True
 257 null other = False
 258
 259 -- | /O(n)/. Number of elements in the map.
 260 size :: IntMap a -> Int
 261 size t
 262   = case t of
 263       Bin p m l r -> size l + size r
 264       Tip k x -> 1
 265       Nil     -> 0
 266
 267 -- | /O(min(n,W))/. Is the key a member of the map?
 268 member :: Key -> IntMap a -> Bool
 269 member k m
 270   = case lookup k m of
 271       Nothing -> False
 272       Just x  -> True
 273
 274 -- | /O(log n)/. Is the key not a member of the map?
 275 notMember :: Key -> IntMap a -> Bool
 276 notMember k m = not $ member k m
 277
 278 -- | /O(min(n,W))/. Lookup the value at a key in the map.
 279 lookup :: (Monad m) => Key -> IntMap a -> m a
 280 lookup k t = case lookup' k t of
 281     Just x -> return x
 282     Nothing -> fail "Data.IntMap.lookup: Key not found"
 283
 284 lookup' :: Key -> IntMap a -> Maybe a
 285 lookup' k t
 286   = let nk = natFromInt k  in seq nk (lookupN nk t)
 287
 288 lookupN :: Nat -> IntMap a -> Maybe a
 289 lookupN k t
 290   = case t of
 291       Bin p m l r
 292         | zeroN k (natFromInt m) -> lookupN k l
 293         | otherwise              -> lookupN k r
 294       Tip kx x
 295         | (k == natFromInt kx)  -> Just x
 296         | otherwise             -> Nothing
 297       Nil -> Nothing
 298
 299 find' :: Key -> IntMap a -> a
 300 find' k m
 301   = case lookup k m of
 302       Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
 303       Just x  -> x
 304
 305
 306 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
 307 -- returns the value at key @k@ or returns @def@ when the key is not an
 308 -- element of the map.
 309 findWithDefault :: a -> Key -> IntMap a -> a
 310 findWithDefault def k m
 311   = case lookup k m of
 312       Nothing -> def
 313       Just x  -> x
 314
 315 {--------------------------------------------------------------------
 316   Construction
 317 --------------------------------------------------------------------}
 318 -- | /O(1)/. The empty map.
 319 empty :: IntMap a
 320 empty
 321   = Nil
 322
 323 -- | /O(1)/. A map of one element.
 324 singleton :: Key -> a -> IntMap a
 325 singleton k x
 326   = Tip k x
 327
 328 {--------------------------------------------------------------------
 329   Insert
 330 --------------------------------------------------------------------}
 331 -- | /O(min(n,W))/. Insert a new key\/value pair in the map.
 332 -- If the key is already present in the map, the associated value is
 333 -- replaced with the supplied value, i.e. 'insert' is equivalent to
 334 -- @'insertWith' 'const'@.
 335 insert :: Key -> a -> IntMap a -> IntMap a
 336 insert k x t
 337   = case t of
 338       Bin p m l r
 339         | nomatch k p m -> join k (Tip k x) p t
 340         | zero k m      -> Bin p m (insert k x l) r
 341         | otherwise     -> Bin p m l (insert k x r)
 342       Tip ky y
 343         | k==ky         -> Tip k x
 344         | otherwise     -> join k (Tip k x) ky t
 345       Nil -> Tip k x
 346
 347 -- right-biased insertion, used by 'union'
 348 -- | /O(min(n,W))/. Insert with a combining function.
 349 -- @'insertWith' f key value mp@
 350 -- will insert the pair (key, value) into @mp@ if key does
 351 -- not exist in the map. If the key does exist, the function will
 352 -- insert @f new_value old_value@.
 353 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
 354 insertWith f k x t
 355   = insertWithKey (\k x y -> f x y) k x t
 356
 357 -- | /O(min(n,W))/. Insert with a combining function.
 358 -- @'insertWithKey' f key value mp@
 359 -- will insert the pair (key, value) into @mp@ if key does
 360 -- not exist in the map. If the key does exist, the function will
 361 -- insert @f key new_value old_value@.
 362 insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
 363 insertWithKey f k x t
 364   = case t of
 365       Bin p m l r
 366         | nomatch k p m -> join k (Tip k x) p t
 367         | zero k m      -> Bin p m (insertWithKey f k x l) r
 368         | otherwise     -> Bin p m l (insertWithKey f k x r)
 369       Tip ky y
 370         | k==ky         -> Tip k (f k x y)
 371         | otherwise     -> join k (Tip k x) ky t
 372       Nil -> Tip k x
 373
 374
 375 -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@)
 376 -- is a pair where the first element is equal to (@'lookup' k map@)
 377 -- and the second element equal to (@'insertWithKey' f k x map@).
 378 insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
 379 insertLookupWithKey f k x t
 380   = case t of
 381       Bin p m l r
 382         | nomatch k p m -> (Nothing,join k (Tip k x) p t)
 383         | zero k m      -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r)
 384         | otherwise     -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r')
 385       Tip ky y
 386         | k==ky         -> (Just y,Tip k (f k x y))
 387         | otherwise     -> (Nothing,join k (Tip k x) ky t)
 388       Nil -> (Nothing,Tip k x)
 389
 390
 391 {--------------------------------------------------------------------
 392   Deletion
 393   [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
 394 --------------------------------------------------------------------}
 395 -- | /O(min(n,W))/. Delete a key and its value from the map. When the key is not
 396 -- a member of the map, the original map is returned.
 397 delete :: Key -> IntMap a -> IntMap a
 398 delete k t
 399   = case t of
 400       Bin p m l r
 401         | nomatch k p m -> t
 402         | zero k m      -> bin p m (delete k l) r
 403         | otherwise     -> bin p m l (delete k r)
 404       Tip ky y
 405         | k==ky         -> Nil
 406         | otherwise     -> t
 407       Nil -> Nil
 408
 409 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
 410 -- a member of the map, the original map is returned.
 411 adjust ::  (a -> a) -> Key -> IntMap a -> IntMap a
 412 adjust f k m
 413   = adjustWithKey (\k x -> f x) k m
 414
 415 -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not
 416 -- a member of the map, the original map is returned.
 417 adjustWithKey ::  (Key -> a -> a) -> Key -> IntMap a -> IntMap a
 418 adjustWithKey f k m
 419   = updateWithKey (\k x -> Just (f k x)) k m
 420
 421 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
 422 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
 423 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
 424 update ::  (a -> Maybe a) -> Key -> IntMap a -> IntMap a
 425 update f k m
 426   = updateWithKey (\k x -> f x) k m
 427
 428 -- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@
 429 -- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
 430 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
 431 updateWithKey ::  (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
 432 updateWithKey f k t
 433   = case t of
 434       Bin p m l r
 435         | nomatch k p m -> t
 436         | zero k m      -> bin p m (updateWithKey f k l) r
 437         | otherwise     -> bin p m l (updateWithKey f k r)
 438       Tip ky y
 439         | k==ky         -> case (f k y) of
 440                              Just y' -> Tip ky y'
 441                              Nothing -> Nil
 442         | otherwise     -> t
 443       Nil -> Nil
 444
 445 -- | /O(min(n,W))/. Lookup and update.
 446 updateLookupWithKey ::  (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
 447 updateLookupWithKey f k t
 448   = case t of
 449       Bin p m l r
 450         | nomatch k p m -> (Nothing,t)
 451         | zero k m      -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r)
 452         | otherwise     -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
 453       Tip ky y
 454         | k==ky         -> case (f k y) of
 455                              Just y' -> (Just y,Tip ky y')
 456                              Nothing -> (Just y,Nil)
 457         | otherwise     -> (Nothing,t)
 458       Nil -> (Nothing,Nil)
 459
 460
 461
 462 -- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
 463 -- 'alter' can be used to insert, delete, or update a value in a 'Map'.
 464 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@
 465 alter f k t
 466   = case t of
 467       Bin p m l r
 468         | nomatch k p m -> case f Nothing of
 469                              Nothing -> t
 470                              Just x -> join k (Tip k x) p t
 471         | zero k m      -> bin p m (alter f k l) r
 472         | otherwise     -> bin p m l (alter f k r)
 473       Tip ky y
 474         | k==ky         -> case f (Just y) of
 475                              Just x -> Tip ky x
 476                              Nothing -> Nil
 477         | otherwise     -> case f Nothing of
 478                              Just x -> join k (Tip k x) ky t
 479                              Nothing -> Tip ky y
 480       Nil               -> case f Nothing of
 481                              Just x -> Tip k x
 482                              Nothing -> Nil
 483
 484
 485 {--------------------------------------------------------------------
 486   Union
 487 --------------------------------------------------------------------}
 488 -- | The union of a list of maps.
 489 unions :: [IntMap a] -> IntMap a
 490 unions xs
 491   = foldlStrict union empty xs
 492
 493 -- | The union of a list of maps, with a combining operation
 494 unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a
 495 unionsWith f ts
 496   = foldlStrict (unionWith f) empty ts
 497
 498 -- | /O(n+m)/. The (left-biased) union of two maps.
 499 -- It prefers the first map when duplicate keys are encountered,
 500 -- i.e. (@'union' == 'unionWith' 'const'@).
 501 union :: IntMap a -> IntMap a -> IntMap a
 502 union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 503   | shorter m1 m2  = union1
 504   | shorter m2 m1  = union2
 505   | p1 == p2       = Bin p1 m1 (union l1 l2) (union r1 r2)
 506   | otherwise      = join p1 t1 p2 t2
 507   where
 508     union1  | nomatch p2 p1 m1  = join p1 t1 p2 t2
 509             | zero p2 m1        = Bin p1 m1 (union l1 t2) r1
 510             | otherwise         = Bin p1 m1 l1 (union r1 t2)
 511
 512     union2  | nomatch p1 p2 m2  = join p1 t1 p2 t2
 513             | zero p1 m2        = Bin p2 m2 (union t1 l2) r2
 514             | otherwise         = Bin p2 m2 l2 (union t1 r2)
 515
 516 union (Tip k x) t = insert k x t
 517 union t (Tip k x) = insertWith (\x y -> y) k x t  -- right bias
 518 union Nil t       = t
 519 union t Nil       = t
 520
 521 -- | /O(n+m)/. The union with a combining function.
 522 unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
 523 unionWith f m1 m2
 524   = unionWithKey (\k x y -> f x y) m1 m2
 525
 526 -- | /O(n+m)/. The union with a combining function.
 527 unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
 528 unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 529   | shorter m1 m2  = union1
 530   | shorter m2 m1  = union2
 531   | p1 == p2       = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
 532   | otherwise      = join p1 t1 p2 t2
 533   where
 534     union1  | nomatch p2 p1 m1  = join p1 t1 p2 t2
 535             | zero p2 m1        = Bin p1 m1 (unionWithKey f l1 t2) r1
 536             | otherwise         = Bin p1 m1 l1 (unionWithKey f r1 t2)
 537
 538     union2  | nomatch p1 p2 m2  = join p1 t1 p2 t2
 539             | zero p1 m2        = Bin p2 m2 (unionWithKey f t1 l2) r2
 540             | otherwise         = Bin p2 m2 l2 (unionWithKey f t1 r2)
 541
 542 unionWithKey f (Tip k x) t = insertWithKey f k x t
 543 unionWithKey f t (Tip k x) = insertWithKey (\k x y -> f k y x) k x t  -- right bias
 544 unionWithKey f Nil t  = t
 545 unionWithKey f t Nil  = t
 546
 547 {--------------------------------------------------------------------
 548   Difference
 549 --------------------------------------------------------------------}
 550 -- | /O(n+m)/. Difference between two maps (based on keys).
 551 difference :: IntMap a -> IntMap b -> IntMap a
 552 difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 553   | shorter m1 m2  = difference1
 554   | shorter m2 m1  = difference2
 555   | p1 == p2       = bin p1 m1 (difference l1 l2) (difference r1 r2)
 556   | otherwise      = t1
 557   where
 558     difference1 | nomatch p2 p1 m1  = t1
 559                 | zero p2 m1        = bin p1 m1 (difference l1 t2) r1
 560                 | otherwise         = bin p1 m1 l1 (difference r1 t2)
 561
 562     difference2 | nomatch p1 p2 m2  = t1
 563                 | zero p1 m2        = difference t1 l2
 564                 | otherwise         = difference t1 r2
 565
 566 difference t1@(Tip k x) t2
 567   | member k t2  = Nil
 568   | otherwise    = t1
 569
 570 difference Nil t       = Nil
 571 difference t (Tip k x) = delete k t
 572 difference t Nil       = t
 573
 574 -- | /O(n+m)/. Difference with a combining function.
 575 differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
 576 differenceWith f m1 m2
 577   = differenceWithKey (\k x y -> f x y) m1 m2
 578
 579 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
 580 -- encountered, the combining function is applied to the key and both values.
 581 -- If it returns 'Nothing', the element is discarded (proper set difference).
 582 -- If it returns (@'Just' y@), the element is updated with a new value @y@.
 583 differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a
 584 differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 585   | shorter m1 m2  = difference1
 586   | shorter m2 m1  = difference2
 587   | p1 == p2       = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2)
 588   | otherwise      = t1
 589   where
 590     difference1 | nomatch p2 p1 m1  = t1
 591                 | zero p2 m1        = bin p1 m1 (differenceWithKey f l1 t2) r1
 592                 | otherwise         = bin p1 m1 l1 (differenceWithKey f r1 t2)
 593
 594     difference2 | nomatch p1 p2 m2  = t1
 595                 | zero p1 m2        = differenceWithKey f t1 l2
 596                 | otherwise         = differenceWithKey f t1 r2
 597
 598 differenceWithKey f t1@(Tip k x) t2
 599   = case lookup k t2 of
 600       Just y  -> case f k x y of
 601                    Just y' -> Tip k y'
 602                    Nothing -> Nil
 603       Nothing -> t1
 604
 605 differenceWithKey f Nil t       = Nil
 606 differenceWithKey f t (Tip k y) = updateWithKey (\k x -> f k x y) k t
 607 differenceWithKey f t Nil       = t
 608
 609
 610 {--------------------------------------------------------------------
 611   Intersection
 612 --------------------------------------------------------------------}
 613 -- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys).
 614 intersection :: IntMap a -> IntMap b -> IntMap a
 615 intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 616   | shorter m1 m2  = intersection1
 617   | shorter m2 m1  = intersection2
 618   | p1 == p2       = bin p1 m1 (intersection l1 l2) (intersection r1 r2)
 619   | otherwise      = Nil
 620   where
 621     intersection1 | nomatch p2 p1 m1  = Nil
 622                   | zero p2 m1        = intersection l1 t2
 623                   | otherwise         = intersection r1 t2
 624
 625     intersection2 | nomatch p1 p2 m2  = Nil
 626                   | zero p1 m2        = intersection t1 l2
 627                   | otherwise         = intersection t1 r2
 628
 629 intersection t1@(Tip k x) t2
 630   | member k t2  = t1
 631   | otherwise    = Nil
 632 intersection t (Tip k x)
 633   = case lookup k t of
 634       Just y  -> Tip k y
 635       Nothing -> Nil
 636 intersection Nil t = Nil
 637 intersection t Nil = Nil
 638
 639 -- | /O(n+m)/. The intersection with a combining function.
 640 intersectionWith :: (a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
 641 intersectionWith f m1 m2
 642   = intersectionWithKey (\k x y -> f x y) m1 m2
 643
 644 -- | /O(n+m)/. The intersection with a combining function.
 645 intersectionWithKey :: (Key -> a -> b -> a) -> IntMap a -> IntMap b -> IntMap a
 646 intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 647   | shorter m1 m2  = intersection1
 648   | shorter m2 m1  = intersection2
 649   | p1 == p2       = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
 650   | otherwise      = Nil
 651   where
 652     intersection1 | nomatch p2 p1 m1  = Nil
 653                   | zero p2 m1        = intersectionWithKey f l1 t2
 654                   | otherwise         = intersectionWithKey f r1 t2
 655
 656     intersection2 | nomatch p1 p2 m2  = Nil
 657                   | zero p1 m2        = intersectionWithKey f t1 l2
 658                   | otherwise         = intersectionWithKey f t1 r2
 659
 660 intersectionWithKey f t1@(Tip k x) t2
 661   = case lookup k t2 of
 662       Just y  -> Tip k (f k x y)
 663       Nothing -> Nil
 664 intersectionWithKey f t1 (Tip k y)
 665   = case lookup k t1 of
 666       Just x  -> Tip k (f k x y)
 667       Nothing -> Nil
 668 intersectionWithKey f Nil t = Nil
 669 intersectionWithKey f t Nil = Nil
 670
 671
 672 {--------------------------------------------------------------------
 673   Submap
 674 --------------------------------------------------------------------}
 675 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
 676 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
 677 isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
 678 isProperSubmapOf m1 m2
 679   = isProperSubmapOfBy (==) m1 m2
 680
 681 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
 682  The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
 683  @m1@ and @m2@ are not equal,
 684  all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
 685  applied to their respective values. For example, the following
 686  expressions are all 'True':
 687
 688   > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 689   > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 690
 691  But the following are all 'False':
 692
 693   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
 694   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
 695   > isProperSubmapOfBy (<)  (fromList [(1,1)])       (fromList [(1,1),(2,2)])
 696 -}
 697 isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
 698 isProperSubmapOfBy pred t1 t2
 699   = case submapCmp pred t1 t2 of
 700       LT -> True
 701       ge -> False
 702
 703 submapCmp pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 704   | shorter m1 m2  = GT
 705   | shorter m2 m1  = submapCmpLt
 706   | p1 == p2       = submapCmpEq
 707   | otherwise      = GT  -- disjoint
 708   where
 709     submapCmpLt | nomatch p1 p2 m2  = GT
 710                 | zero p1 m2        = submapCmp pred t1 l2
 711                 | otherwise         = submapCmp pred t1 r2
 712     submapCmpEq = case (submapCmp pred l1 l2, submapCmp pred r1 r2) of
 713                     (GT,_ ) -> GT
 714                     (_ ,GT) -> GT
 715                     (EQ,EQ) -> EQ
 716                     other   -> LT
 717
 718 submapCmp pred (Bin p m l r) t  = GT
 719 submapCmp pred (Tip kx x) (Tip ky y)
 720   | (kx == ky) && pred x y = EQ
 721   | otherwise              = GT  -- disjoint
 722 submapCmp pred (Tip k x) t
 723   = case lookup k t of
 724      Just y  | pred x y -> LT
 725      other   -> GT -- disjoint
 726 submapCmp pred Nil Nil = EQ
 727 submapCmp pred Nil t   = LT
 728
 729 -- | /O(n+m)/. Is this a submap?
 730 -- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
 731 isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool
 732 isSubmapOf m1 m2
 733   = isSubmapOfBy (==) m1 m2
 734
 735 {- | /O(n+m)/.
 736  The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if
 737  all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
 738  applied to their respective values. For example, the following
 739  expressions are all 'True':
 740
 741   > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 742   > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 743   > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
 744
 745  But the following are all 'False':
 746
 747   > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)])
 748   > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 749   > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
 750 -}
 751
 752 isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool
 753 isSubmapOfBy pred t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2)
 754   | shorter m1 m2  = False
 755   | shorter m2 m1  = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy pred t1 l2
 756                                                       else isSubmapOfBy pred t1 r2)
 757   | otherwise      = (p1==p2) && isSubmapOfBy pred l1 l2 && isSubmapOfBy pred r1 r2
 758 isSubmapOfBy pred (Bin p m l r) t  = False
 759 isSubmapOfBy pred (Tip k x) t      = case lookup k t of
 760                                    Just y  -> pred x y
 761                                    Nothing -> False
 762 isSubmapOfBy pred Nil t            = True
 763
 764 {--------------------------------------------------------------------
 765   Mapping
 766 --------------------------------------------------------------------}
 767 -- | /O(n)/. Map a function over all values in the map.
 768 map :: (a -> b) -> IntMap a -> IntMap b
 769 map f m
 770   = mapWithKey (\k x -> f x) m
 771
 772 -- | /O(n)/. Map a function over all values in the map.
 773 mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b
 774 mapWithKey f t
 775   = case t of
 776       Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r)
 777       Tip k x     -> Tip k (f k x)
 778       Nil         -> Nil
 779
 780 -- | /O(n)/. The function @'mapAccum'@ threads an accumulating
 781 -- argument through the map in ascending order of keys.
 782 mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 783 mapAccum f a m
 784   = mapAccumWithKey (\a k x -> f a x) a m
 785
 786 -- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating
 787 -- argument through the map in ascending order of keys.
 788 mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 789 mapAccumWithKey f a t
 790   = mapAccumL f a t
 791
 792 -- | /O(n)/. The function @'mapAccumL'@ threads an accumulating
 793 -- argument through the map in ascending order of keys.
 794 mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 795 mapAccumL f a t
 796   = case t of
 797       Bin p m l r -> let (a1,l') = mapAccumL f a l
 798                          (a2,r') = mapAccumL f a1 r
 799                      in (a2,Bin p m l' r')
 800       Tip k x     -> let (a',x') = f a k x in (a',Tip k x')
 801       Nil         -> (a,Nil)
 802
 803
 804 -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating
 805 -- argument throught the map in descending order of keys.
 806 mapAccumR :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c)
 807 mapAccumR f a t
 808   = case t of
 809       Bin p m l r -> let (a1,r') = mapAccumR f a r
 810                          (a2,l') = mapAccumR f a1 l
 811                      in (a2,Bin p m l' r')
 812       Tip k x     -> let (a',x') = f a k x in (a',Tip k x')
 813       Nil         -> (a,Nil)
 814
 815 {--------------------------------------------------------------------
 816   Filter
 817 --------------------------------------------------------------------}
 818 -- | /O(n)/. Filter all values that satisfy some predicate.
 819 filter :: (a -> Bool) -> IntMap a -> IntMap a
 820 filter p m
 821   = filterWithKey (\k x -> p x) m
 822
 823 -- | /O(n)/. Filter all keys\/values that satisfy some predicate.
 824 filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a
 825 filterWithKey pred t
 826   = case t of
 827       Bin p m l r
 828         -> bin p m (filterWithKey pred l) (filterWithKey pred r)
 829       Tip k x
 830         | pred k x  -> t
 831         | otherwise -> Nil
 832       Nil -> Nil
 833
 834 -- | /O(n)/. partition the map according to some predicate. The first
 835 -- map contains all elements that satisfy the predicate, the second all
 836 -- elements that fail the predicate. See also 'split'.
 837 partition :: (a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
 838 partition p m
 839   = partitionWithKey (\k x -> p x) m
 840
 841 -- | /O(n)/. partition the map according to some predicate. The first
 842 -- map contains all elements that satisfy the predicate, the second all
 843 -- elements that fail the predicate. See also 'split'.
 844 partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a)
 845 partitionWithKey pred t
 846   = case t of
 847       Bin p m l r
 848         -> let (l1,l2) = partitionWithKey pred l
 849                (r1,r2) = partitionWithKey pred r
 850            in (bin p m l1 r1, bin p m l2 r2)
 851       Tip k x
 852         | pred k x  -> (t,Nil)
 853         | otherwise -> (Nil,t)
 854       Nil -> (Nil,Nil)
 855
 856 -- | /O(n)/. Map values and collect the 'Just' results.
 857 mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b
 858 mapMaybe f m
 859   = mapMaybeWithKey (\k x -> f x) m
 860
 861 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
 862 mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b
 863 mapMaybeWithKey f (Bin p m l r)
 864   = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r)
 865 mapMaybeWithKey f (Tip k x) = case f k x of
 866   Just y  -> Tip k y
 867   Nothing -> Nil
 868 mapMaybeWithKey f Nil = Nil
 869
 870 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
 871 mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
 872 mapEither f m
 873   = mapEitherWithKey (\k x -> f x) m
 874
 875 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
 876 mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
 877 mapEitherWithKey f (Bin p m l r)
 878   = (bin p m l1 r1, bin p m l2 r2)
 879   where
 880     (l1,l2) = mapEitherWithKey f l
 881     (r1,r2) = mapEitherWithKey f r
 882 mapEitherWithKey f (Tip k x) = case f k x of
 883   Left y  -> (Tip k y, Nil)
 884   Right z -> (Nil, Tip k z)
 885 mapEitherWithKey f Nil = (Nil, Nil)
 886
 887 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@
 888 -- where all keys in @map1@ are lower than @k@ and all keys in
 889 -- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
 890 split :: Key -> IntMap a -> (IntMap a,IntMap a)
 891 split k t
 892   = case t of
 893       Bin p m l r
 894           | m < 0 -> (if k >= 0 -- handle negative numbers.
 895                       then let (lt,gt) = split' k l in (union r lt, gt)
 896                       else let (lt,gt) = split' k r in (lt, union gt l))
 897           | otherwise   -> split' k t
 898       Tip ky y
 899         | k>ky      -> (t,Nil)
 900         | k<ky      -> (Nil,t)
 901         | otherwise -> (Nil,Nil)
 902       Nil -> (Nil,Nil)
 903
 904 split' :: Key -> IntMap a -> (IntMap a,IntMap a)
 905 split' k t
 906   = case t of
 907       Bin p m l r
 908         | nomatch k p m -> if k>p then (t,Nil) else (Nil,t)
 909         | zero k m  -> let (lt,gt) = split k l in (lt,union gt r)
 910         | otherwise -> let (lt,gt) = split k r in (union l lt,gt)
 911       Tip ky y
 912         | k>ky      -> (t,Nil)
 913         | k<ky      -> (Nil,t)
 914         | otherwise -> (Nil,Nil)
 915       Nil -> (Nil,Nil)
 916
 917 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
 918 -- key was found in the original map.
 919 splitLookup :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
 920 splitLookup k t
 921   = case t of
 922       Bin p m l r
 923           | m < 0 -> (if k >= 0 -- handle negative numbers.
 924                       then let (lt,found,gt) = splitLookup' k l in (union r lt,found, gt)
 925                       else let (lt,found,gt) = splitLookup' k r in (lt,found, union gt l))
 926           | otherwise   -> splitLookup' k t
 927       Tip ky y
 928         | k>ky      -> (t,Nothing,Nil)
 929         | k<ky      -> (Nil,Nothing,t)
 930         | otherwise -> (Nil,Just y,Nil)
 931       Nil -> (Nil,Nothing,Nil)
 932
 933 splitLookup' :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a)
 934 splitLookup' k t
 935   = case t of
 936       Bin p m l r
 937         | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t)
 938         | zero k m  -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r)
 939         | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt)
 940       Tip ky y
 941         | k>ky      -> (t,Nothing,Nil)
 942         | k<ky      -> (Nil,Nothing,t)
 943         | otherwise -> (Nil,Just y,Nil)
 944       Nil -> (Nil,Nothing,Nil)
 945
 946 {--------------------------------------------------------------------
 947   Fold
 948 --------------------------------------------------------------------}
 949 -- | /O(n)/. Fold the values in the map, such that
 950 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
 951 -- For example,
 952 --
 953 -- > elems map = fold (:) [] map
 954 --
 955 fold :: (a -> b -> b) -> b -> IntMap a -> b
 956 fold f z t
 957   = foldWithKey (\k x y -> f x y) z t
 958
 959 -- | /O(n)/. Fold the keys and values in the map, such that
 960 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
 961 -- For example,
 962 --
 963 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
 964 --
 965 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
 966 foldWithKey f z t
 967   = foldr f z t
 968
 969 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
 970 foldr f z t
 971   = case t of
 972       Bin 0 m l r | m < 0 -> foldr' f (foldr' f z l) r  -- put negative numbers before.
 973       Bin _ _ _ _ -> foldr' f z t
 974       Tip k x     -> f k x z
 975       Nil         -> z
 976
 977 foldr' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
 978 foldr' f z t
 979   = case t of
 980       Bin p m l r -> foldr' f (foldr' f z r) l
 981       Tip k x     -> f k x z
 982       Nil         -> z
 983
 984
 985
 986 {--------------------------------------------------------------------
 987   List variations
 988 --------------------------------------------------------------------}
 989 -- | /O(n)/.
 990 -- Return all elements of the map in the ascending order of their keys.
 991 elems :: IntMap a -> [a]
 992 elems m
 993   = foldWithKey (\k x xs -> x:xs) [] m
 994
 995 -- | /O(n)/. Return all keys of the map in ascending order.
 996 keys  :: IntMap a -> [Key]
 997 keys m
 998   = foldWithKey (\k x ks -> k:ks) [] m
 999
1000 -- | /O(n*min(n,W))/. The set of all keys of the map.
1001 keysSet :: IntMap a -> IntSet.IntSet
1002 keysSet m = IntSet.fromDistinctAscList (keys m)
1003
1004
1005 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
1006 assocs :: IntMap a -> [(Key,a)]
1007 assocs m
1008   = toList m
1009
1010
1011 {--------------------------------------------------------------------
1012   Lists
1013 --------------------------------------------------------------------}
1014 -- | /O(n)/. Convert the map to a list of key\/value pairs.
1015 toList :: IntMap a -> [(Key,a)]
1016 toList t
1017   = foldWithKey (\k x xs -> (k,x):xs) [] t
1018
1019 -- | /O(n)/. Convert the map to a list of key\/value pairs where the
1020 -- keys are in ascending order.
1021 toAscList :: IntMap a -> [(Key,a)]
1022 toAscList t
1023   = -- NOTE: the following algorithm only works for big-endian trees
1024     let (pos,neg) = span (\(k,x) -> k >=0) (foldr (\k x xs -> (k,x):xs) [] t) in neg ++ pos
1025
1026 -- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs.
1027 fromList :: [(Key,a)] -> IntMap a
1028 fromList xs
1029   = foldlStrict ins empty xs
1030   where
1031     ins t (k,x)  = insert k x t
1032
1033 -- | /O(n*min(n,W))/.  Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1034 fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1035 fromListWith f xs
1036   = fromListWithKey (\k x y -> f x y) xs
1037
1038 -- | /O(n*min(n,W))/.  Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
1039 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1040 fromListWithKey f xs
1041   = foldlStrict ins empty xs
1042   where
1043     ins t (k,x) = insertWithKey f k x t
1044
1045 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1046 -- the keys are in ascending order.
1047 fromAscList :: [(Key,a)] -> IntMap a
1048 fromAscList xs
1049   = fromList xs
1050
1051 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1052 -- the keys are in ascending order, with a combining function on equal keys.
1053 fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
1054 fromAscListWith f xs
1055   = fromListWith f xs
1056
1057 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1058 -- the keys are in ascending order, with a combining function on equal keys.
1059 fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
1060 fromAscListWithKey f xs
1061   = fromListWithKey f xs
1062
1063 -- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs where
1064 -- the keys are in ascending order and all distinct.
1065 fromDistinctAscList :: [(Key,a)] -> IntMap a
1066 fromDistinctAscList xs
1067   = fromList xs
1068
1069
1070 {--------------------------------------------------------------------
1071   Eq
1072 --------------------------------------------------------------------}
1073 instance Eq a => Eq (IntMap a) where
1074   t1 == t2  = equal t1 t2
1075   t1 /= t2  = nequal t1 t2
1076
1077 equal :: Eq a => IntMap a -> IntMap a -> Bool
1078 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1079   = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
1080 equal (Tip kx x) (Tip ky y)
1081   = (kx == ky) && (x==y)
1082 equal Nil Nil = True
1083 equal t1 t2   = False
1084
1085 nequal :: Eq a => IntMap a -> IntMap a -> Bool
1086 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
1087   = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
1088 nequal (Tip kx x) (Tip ky y)
1089   = (kx /= ky) || (x/=y)
1090 nequal Nil Nil = False
1091 nequal t1 t2   = True
1092
1093 {--------------------------------------------------------------------
1094   Ord
1095 --------------------------------------------------------------------}
1096
1097 instance Ord a => Ord (IntMap a) where
1098     compare m1 m2 = compare (toList m1) (toList m2)
1099
1100 {--------------------------------------------------------------------
1101   Functor
1102 --------------------------------------------------------------------}
1103
1104 instance Functor IntMap where
1105     fmap = map
1106
1107 {--------------------------------------------------------------------
1108   Show
1109 --------------------------------------------------------------------}
1110
1111 instance Show a => Show (IntMap a) where
1112   showsPrec d m   = showParen (d > 10) $
1113     showString "fromList " . shows (toList m)
1114
1115 showMap :: (Show a) => [(Key,a)] -> ShowS
1116 showMap []
1117   = showString "{}"
1118 showMap (x:xs)
1119   = showChar '{' . showElem x . showTail xs
1120   where
1121     showTail []     = showChar '}'
1122     showTail (x:xs) = showChar ',' . showElem x . showTail xs
1123
1124     showElem (k,x)  = shows k . showString ":=" . shows x
1125
1126 {--------------------------------------------------------------------
1127   Read
1128 --------------------------------------------------------------------}
1129 instance (Read e) => Read (IntMap e) where
1130 #ifdef __GLASGOW_HASKELL__
1131   readPrec = parens $ prec 10 $ do
1132     Ident "fromList" <- lexP
1133     xs <- readPrec
1134     return (fromList xs)
1135
1136   readListPrec = readListPrecDefault
1137 #else
1138   readsPrec p = readParen (p > 10) $ \ r -> do
1139     ("fromList",s) <- lex r
1140     (xs,t) <- reads s
1141     return (fromList xs,t)
1142 #endif
1143
1144 {--------------------------------------------------------------------
1145   Typeable
1146 --------------------------------------------------------------------}
1147
1148 #include "Typeable.h"
1149 INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
1150
1151 {--------------------------------------------------------------------
1152   Debugging
1153 --------------------------------------------------------------------}
1154 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1155 -- in a compressed, hanging format.
1156 showTree :: Show a => IntMap a -> String
1157 showTree s
1158   = showTreeWith True False s
1159
1160
1161 {- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows
1162  the tree that implements the map. If @hang@ is
1163  'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1164  @wide@ is 'True', an extra wide version is shown.
1165 -}
1166 showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String
1167 showTreeWith hang wide t
1168   | hang      = (showsTreeHang wide [] t) ""
1169   | otherwise = (showsTree wide [] [] t) ""
1170
1171 showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS
1172 showsTree wide lbars rbars t
1173   = case t of
1174       Bin p m l r
1175           -> showsTree wide (withBar rbars) (withEmpty rbars) r .
1176              showWide wide rbars .
1177              showsBars lbars . showString (showBin p m) . showString "\n" .
1178              showWide wide lbars .
1179              showsTree wide (withEmpty lbars) (withBar lbars) l
1180       Tip k x
1181           -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1182       Nil -> showsBars lbars . showString "|\n"
1183
1184 showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS
1185 showsTreeHang wide bars t
1186   = case t of
1187       Bin p m l r
1188           -> showsBars bars . showString (showBin p m) . showString "\n" .
1189              showWide wide bars .
1190              showsTreeHang wide (withBar bars) l .
1191              showWide wide bars .
1192              showsTreeHang wide (withEmpty bars) r
1193       Tip k x
1194           -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n"
1195       Nil -> showsBars bars . showString "|\n"
1196
1197 showBin p m
1198   = "*" -- ++ show (p,m)
1199
1200 showWide wide bars
1201   | wide      = showString (concat (reverse bars)) . showString "|\n"
1202   | otherwise = id
1203
1204 showsBars :: [String] -> ShowS
1205 showsBars bars
1206   = case bars of
1207       [] -> id
1208       _  -> showString (concat (reverse (tail bars))) . showString node
1209
1210 node           = "+--"
1211 withBar bars   = "|  ":bars
1212 withEmpty bars = "   ":bars
1213
1214
1215 {--------------------------------------------------------------------
1216   Helpers
1217 --------------------------------------------------------------------}
1218 {--------------------------------------------------------------------
1219   Join
1220 --------------------------------------------------------------------}
1221 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
1222 join p1 t1 p2 t2
1223   | zero p1 m = Bin p m t1 t2
1224   | otherwise = Bin p m t2 t1
1225   where
1226     m = branchMask p1 p2
1227     p = mask p1 m
1228
1229 {--------------------------------------------------------------------
1230   @bin@ assures that we never have empty trees within a tree.
1231 --------------------------------------------------------------------}
1232 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
1233 bin p m l Nil = l
1234 bin p m Nil r = r
1235 bin p m l r   = Bin p m l r
1236
1237
1238 {--------------------------------------------------------------------
1239   Endian independent bit twiddling
1240 --------------------------------------------------------------------}
1241 zero :: Key -> Mask -> Bool
1242 zero i m
1243   = (natFromInt i) .&. (natFromInt m) == 0
1244
1245 nomatch,match :: Key -> Prefix -> Mask -> Bool
1246 nomatch i p m
1247   = (mask i m) /= p
1248
1249 match i p m
1250   = (mask i m) == p
1251
1252 mask :: Key -> Mask -> Prefix
1253 mask i m
1254   = maskW (natFromInt i) (natFromInt m)
1255
1256
1257 zeroN :: Nat -> Nat -> Bool
1258 zeroN i m = (i .&. m) == 0
1259
1260 {--------------------------------------------------------------------
1261   Big endian operations
1262 --------------------------------------------------------------------}
1263 maskW :: Nat -> Nat -> Prefix
1264 maskW i m
1265   = intFromNat (i .&. (complement (m-1) `xor` m))
1266
1267 shorter :: Mask -> Mask -> Bool
1268 shorter m1 m2
1269   = (natFromInt m1) > (natFromInt m2)
1270
1271 branchMask :: Prefix -> Prefix -> Mask
1272 branchMask p1 p2
1273   = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
1274
1275 {----------------------------------------------------------------------
1276   Finding the highest bit (mask) in a word [x] can be done efficiently in
1277   three ways:
1278   * convert to a floating point value and the mantissa tells us the
1279     [log2(x)] that corresponds with the highest bit position. The mantissa
1280     is retrieved either via the standard C function [frexp] or by some bit
1281     twiddling on IEEE compatible numbers (float). Note that one needs to
1282     use at least [double] precision for an accurate mantissa of 32 bit
1283     numbers.
1284   * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
1285   * use processor specific assembler instruction (asm).
1286
1287   The most portable way would be [bit], but is it efficient enough?
1288   I have measured the cycle counts of the different methods on an AMD
1289   Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
1290
1291   highestBitMask: method  cycles
1292                   --------------
1293                    frexp   200
1294                    float    33
1295                    bit      11
1296                    asm      12
1297
1298   highestBit:     method  cycles
1299                   --------------
1300                    frexp   195
1301                    float    33
1302                    bit      11
1303                    asm      11
1304
1305   Wow, the bit twiddling is on today's RISC like machines even faster
1306   than a single CISC instruction (BSR)!
1307 ----------------------------------------------------------------------}
1308
1309 {----------------------------------------------------------------------
1310   [highestBitMask] returns a word where only the highest bit is set.
1311   It is found by first setting all bits in lower positions than the
1312   highest bit and than taking an exclusive or with the original value.
1313   Allthough the function may look expensive, GHC compiles this into
1314   excellent C code that subsequently compiled into highly efficient
1315   machine code. The algorithm is derived from Jorg Arndt's FXT library.
1316 ----------------------------------------------------------------------}
1317 highestBitMask :: Nat -> Nat
1318 highestBitMask x
1319   = case (x .|. shiftRL x 1) of
1320      x -> case (x .|. shiftRL x 2) of
1321       x -> case (x .|. shiftRL x 4) of
1322        x -> case (x .|. shiftRL x 8) of
1323         x -> case (x .|. shiftRL x 16) of
1324          x -> case (x .|. shiftRL x 32) of   -- for 64 bit platforms
1325           x -> (x `xor` (shiftRL x 1))
1326
1327
1328 {--------------------------------------------------------------------
1329   Utilities
1330 --------------------------------------------------------------------}
1331 foldlStrict f z xs
1332   = case xs of
1333       []     -> z
1334       (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1335
1336 {-
1337 {--------------------------------------------------------------------
1338   Testing
1339 --------------------------------------------------------------------}
1340 testTree :: [Int] -> IntMap Int
1341 testTree xs   = fromList [(x,x*x*30696 `mod` 65521) | x <- xs]
1342 test1 = testTree [1..20]
1343 test2 = testTree [30,29..10]
1344 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1345
1346 {--------------------------------------------------------------------
1347   QuickCheck
1348 --------------------------------------------------------------------}
1349 qcheck prop
1350   = check config prop
1351   where
1352     config = Config
1353       { configMaxTest = 500
1354       , configMaxFail = 5000
1355       , configSize    = \n -> (div n 2 + 3)
1356       , configEvery   = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1357       }
1358
1359
1360 {--------------------------------------------------------------------
1361   Arbitrary, reasonably balanced trees
1362 --------------------------------------------------------------------}
1363 instance Arbitrary a => Arbitrary (IntMap a) where
1364   arbitrary = do{ ks <- arbitrary
1365                 ; xs <- mapM (\k -> do{ x <- arbitrary; return (k,x)}) ks
1366                 ; return (fromList xs)
1367                 }
1368
1369
1370 {--------------------------------------------------------------------
1371   Single, Insert, Delete
1372 --------------------------------------------------------------------}
1373 prop_Single :: Key -> Int -> Bool
1374 prop_Single k x
1375   = (insert k x empty == singleton k x)
1376
1377 prop_InsertDelete :: Key -> Int -> IntMap Int -> Property
1378 prop_InsertDelete k x t
1379   = not (member k t) ==> delete k (insert k x t) == t
1380
1381 prop_UpdateDelete :: Key -> IntMap Int -> Bool
1382 prop_UpdateDelete k t
1383   = update (const Nothing) k t == delete k t
1384
1385
1386 {--------------------------------------------------------------------
1387   Union
1388 --------------------------------------------------------------------}
1389 prop_UnionInsert :: Key -> Int -> IntMap Int -> Bool
1390 prop_UnionInsert k x t
1391   = union (singleton k x) t == insert k x t
1392
1393 prop_UnionAssoc :: IntMap Int -> IntMap Int -> IntMap Int -> Bool
1394 prop_UnionAssoc t1 t2 t3
1395   = union t1 (union t2 t3) == union (union t1 t2) t3
1396
1397 prop_UnionComm :: IntMap Int -> IntMap Int -> Bool
1398 prop_UnionComm t1 t2
1399   = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1400
1401
1402 prop_Diff :: [(Key,Int)] -> [(Key,Int)] -> Bool
1403 prop_Diff xs ys
1404   =  List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1405     == List.sort ((List.\\) (nub (Prelude.map fst xs))  (nub (Prelude.map fst ys)))
1406
1407 prop_Int :: [(Key,Int)] -> [(Key,Int)] -> Bool
1408 prop_Int xs ys
1409   =  List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1410     == List.sort (nub ((List.intersect) (Prelude.map fst xs)  (Prelude.map fst ys)))
1411
1412 {--------------------------------------------------------------------
1413   Lists
1414 --------------------------------------------------------------------}
1415 prop_Ordered
1416   = forAll (choose (5,100)) $ \n ->
1417     let xs = [(x,()) | x <- [0..n::Int]]
1418     in fromAscList xs == fromList xs
1419
1420 prop_List :: [Key] -> Bool
1421 prop_List xs
1422   = (sort (nub xs) == [x | (x,()) <- toAscList (fromList [(x,()) | x <- xs])])
1423 -}