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