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