Data/IntMap.hs

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