Data/IntMap.hs

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