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