Data/IntMap.hs

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