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