Data/IntMap.hs

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