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