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