Data/IntMap.hs

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