Data/Map.hs

   1 -----------------------------------------------------------------------------
   2 -- |
   3 -- Module      :  Data.Map
   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 keys to values (dictionaries).
  11 --
  12 -- This module is intended to be imported @qualified@, to avoid name
  13 -- clashes with Prelude functions.  eg.
  14 --
  15 -- >  import Data.Map as Map
  16 --
  17 -- The implementation of 'Map' is based on /size balanced/ binary trees (or
  18 -- trees of /bounded balance/) as described by:
  19 --
  20 --    * Stephen Adams, \"/Efficient sets: a balancing act/\",
  21 --      Journal of Functional Programming 3(4):553-562, October 1993,
  22 --      <http://www.swiss.ai.mit.edu/~adams/BB>.
  23 --
  24 --    * J. Nievergelt and E.M. Reingold,
  25 --      \"/Binary search trees of bounded balance/\",
  26 --      SIAM journal of computing 2(1), March 1973.
  27 -----------------------------------------------------------------------------
  28
  29 module Data.Map  (
  30             -- * Map type
  31               Map          -- instance Eq,Show,Read
  32
  33             -- * Operators
  34             , (!), (\\)
  35
  36
  37             -- * Query
  38             , null
  39             , size
  40             , member
  41             , lookup
  42             , findWithDefault
  43
  44             -- * Construction
  45             , empty
  46             , singleton
  47
  48             -- ** Insertion
  49             , insert
  50             , insertWith, insertWithKey, insertLookupWithKey
  51
  52             -- ** Delete\/Update
  53             , delete
  54             , adjust
  55             , adjustWithKey
  56             , update
  57             , updateWithKey
  58             , updateLookupWithKey
  59
  60             -- * Combine
  61
  62             -- ** Union
  63             , union
  64             , unionWith
  65             , unionWithKey
  66             , unions
  67             , unionsWith
  68
  69             -- ** Difference
  70             , difference
  71             , differenceWith
  72             , differenceWithKey
  73
  74             -- ** Intersection
  75             , intersection
  76             , intersectionWith
  77             , intersectionWithKey
  78
  79             -- * Traversal
  80             -- ** Map
  81             , map
  82             , mapWithKey
  83             , mapAccum
  84             , mapAccumWithKey
  85             , mapKeys
  86             , mapKeysWith
  87             , mapKeysMonotonic
  88
  89             -- ** Fold
  90             , fold
  91             , foldWithKey
  92
  93             -- * Conversion
  94             , elems
  95             , keys
  96             , keysSet
  97             , assocs
  98
  99             -- ** Lists
 100             , toList
 101             , fromList
 102             , fromListWith
 103             , fromListWithKey
 104
 105             -- ** Ordered lists
 106             , toAscList
 107             , fromAscList
 108             , fromAscListWith
 109             , fromAscListWithKey
 110             , fromDistinctAscList
 111
 112             -- * Filter
 113             , filter
 114             , filterWithKey
 115             , partition
 116             , partitionWithKey
 117
 118             , split
 119             , splitLookup
 120
 121             -- * Submap
 122             , isSubmapOf, isSubmapOfBy
 123             , isProperSubmapOf, isProperSubmapOfBy
 124
 125             -- * Indexed
 126             , lookupIndex
 127             , findIndex
 128             , elemAt
 129             , updateAt
 130             , deleteAt
 131
 132             -- * Min\/Max
 133             , findMin
 134             , findMax
 135             , deleteMin
 136             , deleteMax
 137             , deleteFindMin
 138             , deleteFindMax
 139             , updateMin
 140             , updateMax
 141             , updateMinWithKey
 142             , updateMaxWithKey
 143
 144             -- * Debugging
 145             , showTree
 146             , showTreeWith
 147             , valid
 148             ) where
 149
 150 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
 151 import qualified Data.Set as Set
 152 import qualified Data.List as List
 153 import Data.Monoid (Monoid(..))
 154 import Data.Typeable
 155
 156 {-
 157 -- for quick check
 158 import qualified Prelude
 159 import qualified List
 160 import Debug.QuickCheck
 161 import List(nub,sort)
 162 -}
 163
 164 #if __GLASGOW_HASKELL__
 165 import Text.Read
 166 import Data.Generics.Basics
 167 import Data.Generics.Instances
 168 #endif
 169
 170 {--------------------------------------------------------------------
 171   Operators
 172 --------------------------------------------------------------------}
 173 infixl 9 !,\\ --
 174
 175 -- | /O(log n)/. Find the value at a key.
 176 -- Calls 'error' when the element can not be found.
 177 (!) :: Ord k => Map k a -> k -> a
 178 m ! k    = find k m
 179
 180 -- | /O(n+m)/. See 'difference'.
 181 (\\) :: Ord k => Map k a -> Map k b -> Map k a
 182 m1 \\ m2 = difference m1 m2
 183
 184 {--------------------------------------------------------------------
 185   Size balanced trees.
 186 --------------------------------------------------------------------}
 187 -- | A Map from keys @k@ to values @a@.
 188 data Map k a  = Tip
 189               | Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a)
 190
 191 type Size     = Int
 192
 193 instance (Ord k) => Monoid (Map k v) where
 194     mempty  = empty
 195     mappend = union
 196     mconcat = unions
 197
 198 #if __GLASGOW_HASKELL__
 199
 200 {--------------------------------------------------------------------
 201   A Data instance
 202 --------------------------------------------------------------------}
 203
 204 -- This instance preserves data abstraction at the cost of inefficiency.
 205 -- We omit reflection services for the sake of data abstraction.
 206
 207 instance (Data k, Data a, Ord k) => Data (Map k a) where
 208   gfoldl f z map = z fromList `f` (toList map)
 209   toConstr _     = error "toConstr"
 210   gunfold _ _    = error "gunfold"
 211   dataTypeOf _   = mkNorepType "Data.Map.Map"
 212
 213 #endif
 214
 215 {--------------------------------------------------------------------
 216   Query
 217 --------------------------------------------------------------------}
 218 -- | /O(1)/. Is the map empty?
 219 null :: Map k a -> Bool
 220 null t
 221   = case t of
 222       Tip             -> True
 223       Bin sz k x l r  -> False
 224
 225 -- | /O(1)/. The number of elements in the map.
 226 size :: Map k a -> Int
 227 size t
 228   = case t of
 229       Tip             -> 0
 230       Bin sz k x l r  -> sz
 231
 232
 233 -- | /O(log n)/. Lookup the value at a key in the map.
 234 lookup :: (Monad m,Ord k) => k -> Map k a -> m a
 235 lookup k t = case lookup' k t of
 236     Just x -> return x
 237     Nothing -> fail "Data.Map.lookup: Key not found"
 238 lookup' :: Ord k => k -> Map k a -> Maybe a
 239 lookup' k t
 240   = case t of
 241       Tip -> Nothing
 242       Bin sz kx x l r
 243           -> case compare k kx of
 244                LT -> lookup' k l
 245                GT -> lookup' k r
 246                EQ -> Just x
 247
 248 -- | /O(log n)/. Is the key a member of the map?
 249 member :: Ord k => k -> Map k a -> Bool
 250 member k m
 251   = case lookup k m of
 252       Nothing -> False
 253       Just x  -> True
 254
 255 -- | /O(log n)/. Find the value at a key.
 256 -- Calls 'error' when the element can not be found.
 257 find :: Ord k => k -> Map k a -> a
 258 find k m
 259   = case lookup k m of
 260       Nothing -> error "Map.find: element not in the map"
 261       Just x  -> x
 262
 263 -- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns
 264 -- the value at key @k@ or returns @def@ when the key is not in the map.
 265 findWithDefault :: Ord k => a -> k -> Map k a -> a
 266 findWithDefault def k m
 267   = case lookup k m of
 268       Nothing -> def
 269       Just x  -> x
 270
 271
 272
 273 {--------------------------------------------------------------------
 274   Construction
 275 --------------------------------------------------------------------}
 276 -- | /O(1)/. The empty map.
 277 empty :: Map k a
 278 empty
 279   = Tip
 280
 281 -- | /O(1)/. A map with a single element.
 282 singleton :: k -> a -> Map k a
 283 singleton k x
 284   = Bin 1 k x Tip Tip
 285
 286 {--------------------------------------------------------------------
 287   Insertion
 288 --------------------------------------------------------------------}
 289 -- | /O(log n)/. Insert a new key and value in the map.
 290 -- If the key is already present in the map, the associated value is
 291 -- replaced with the supplied value, i.e. 'insert' is equivalent to
 292 -- @'insertWith' 'const'@.
 293 insert :: Ord k => k -> a -> Map k a -> Map k a
 294 insert kx x t
 295   = case t of
 296       Tip -> singleton kx x
 297       Bin sz ky y l r
 298           -> case compare kx ky of
 299                LT -> balance ky y (insert kx x l) r
 300                GT -> balance ky y l (insert kx x r)
 301                EQ -> Bin sz kx x l r
 302
 303 -- | /O(log n)/. Insert with a combining function.
 304 -- @'insertWith' f key value mp@
 305 -- will insert the pair (key, value) into @mp@ if key does
 306 -- not exist in the map. If the key does exist, the function will
 307 -- insert @f new_value old_value@.
 308 insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
 309 insertWith f k x m
 310   = insertWithKey (\k x y -> f x y) k x m
 311
 312 -- | /O(log n)/. Insert with a combining function.
 313 -- @'insertWithKey' f key value mp@
 314 -- will insert the pair (key, value) into @mp@ if key does
 315 -- not exist in the map. If the key does exist, the function will
 316 -- insert @f key new_value old_value@.
 317 insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
 318 insertWithKey f kx x t
 319   = case t of
 320       Tip -> singleton kx x
 321       Bin sy ky y l r
 322           -> case compare kx ky of
 323                LT -> balance ky y (insertWithKey f kx x l) r
 324                GT -> balance ky y l (insertWithKey f kx x r)
 325                EQ -> Bin sy ky (f ky x y) l r
 326
 327 -- | /O(log n)/. The expression (@'insertLookupWithKey' f k x map@)
 328 -- is a pair where the first element is equal to (@'lookup' k map@)
 329 -- and the second element equal to (@'insertWithKey' f k x map@).
 330 insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a,Map k a)
 331 insertLookupWithKey f kx x t
 332   = case t of
 333       Tip -> (Nothing, singleton kx x)
 334       Bin sy ky y l r
 335           -> case compare kx ky of
 336                LT -> let (found,l') = insertLookupWithKey f kx x l in (found,balance ky y l' r)
 337                GT -> let (found,r') = insertLookupWithKey f kx x r in (found,balance ky y l r')
 338                EQ -> (Just y, Bin sy ky (f ky x y) l r)
 339
 340 {--------------------------------------------------------------------
 341   Deletion
 342   [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
 343 --------------------------------------------------------------------}
 344 -- | /O(log n)/. Delete a key and its value from the map. When the key is not
 345 -- a member of the map, the original map is returned.
 346 delete :: Ord k => k -> Map k a -> Map k a
 347 delete k t
 348   = case t of
 349       Tip -> Tip
 350       Bin sx kx x l r
 351           -> case compare k kx of
 352                LT -> balance kx x (delete k l) r
 353                GT -> balance kx x l (delete k r)
 354                EQ -> glue l r
 355
 356 -- | /O(log n)/. Adjust a value at a specific key. When the key is not
 357 -- a member of the map, the original map is returned.
 358 adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
 359 adjust f k m
 360   = adjustWithKey (\k x -> f x) k m
 361
 362 -- | /O(log n)/. Adjust a value at a specific key. When the key is not
 363 -- a member of the map, the original map is returned.
 364 adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
 365 adjustWithKey f k m
 366   = updateWithKey (\k x -> Just (f k x)) k m
 367
 368 -- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@
 369 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
 370 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
 371 update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
 372 update f k m
 373   = updateWithKey (\k x -> f x) k m
 374
 375 -- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the
 376 -- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing',
 377 -- the element is deleted. If it is (@'Just' y@), the key @k@ is bound
 378 -- to the new value @y@.
 379 updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
 380 updateWithKey f k t
 381   = case t of
 382       Tip -> Tip
 383       Bin sx kx x l r
 384           -> case compare k kx of
 385                LT -> balance kx x (updateWithKey f k l) r
 386                GT -> balance kx x l (updateWithKey f k r)
 387                EQ -> case f kx x of
 388                        Just x' -> Bin sx kx x' l r
 389                        Nothing -> glue l r
 390
 391 -- | /O(log n)/. Lookup and update.
 392 updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a)
 393 updateLookupWithKey f k t
 394   = case t of
 395       Tip -> (Nothing,Tip)
 396       Bin sx kx x l r
 397           -> case compare k kx of
 398                LT -> let (found,l') = updateLookupWithKey f k l in (found,balance kx x l' r)
 399                GT -> let (found,r') = updateLookupWithKey f k r in (found,balance kx x l r')
 400                EQ -> case f kx x of
 401                        Just x' -> (Just x',Bin sx kx x' l r)
 402                        Nothing -> (Just x,glue l r)
 403
 404 {--------------------------------------------------------------------
 405   Indexing
 406 --------------------------------------------------------------------}
 407 -- | /O(log n)/. Return the /index/ of a key. The index is a number from
 408 -- /0/ up to, but not including, the 'size' of the map. Calls 'error' when
 409 -- the key is not a 'member' of the map.
 410 findIndex :: Ord k => k -> Map k a -> Int
 411 findIndex k t
 412   = case lookupIndex k t of
 413       Nothing  -> error "Map.findIndex: element is not in the map"
 414       Just idx -> idx
 415
 416 -- | /O(log n)/. Lookup the /index/ of a key. The index is a number from
 417 -- /0/ up to, but not including, the 'size' of the map.
 418 lookupIndex :: (Monad m,Ord k) => k -> Map k a -> m Int
 419 lookupIndex k t = case lookup 0 t of
 420     Nothing -> fail "Data.Map.lookupIndex: Key not found."
 421     Just x -> return x
 422   where
 423     lookup idx Tip  = Nothing
 424     lookup idx (Bin _ kx x l r)
 425       = case compare k kx of
 426           LT -> lookup idx l
 427           GT -> lookup (idx + size l + 1) r
 428           EQ -> Just (idx + size l)
 429
 430 -- | /O(log n)/. Retrieve an element by /index/. Calls 'error' when an
 431 -- invalid index is used.
 432 elemAt :: Int -> Map k a -> (k,a)
 433 elemAt i Tip = error "Map.elemAt: index out of range"
 434 elemAt i (Bin _ kx x l r)
 435   = case compare i sizeL of
 436       LT -> elemAt i l
 437       GT -> elemAt (i-sizeL-1) r
 438       EQ -> (kx,x)
 439   where
 440     sizeL = size l
 441
 442 -- | /O(log n)/. Update the element at /index/. Calls 'error' when an
 443 -- invalid index is used.
 444 updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
 445 updateAt f i Tip  = error "Map.updateAt: index out of range"
 446 updateAt f i (Bin sx kx x l r)
 447   = case compare i sizeL of
 448       LT -> updateAt f i l
 449       GT -> updateAt f (i-sizeL-1) r
 450       EQ -> case f kx x of
 451               Just x' -> Bin sx kx x' l r
 452               Nothing -> glue l r
 453   where
 454     sizeL = size l
 455
 456 -- | /O(log n)/. Delete the element at /index/.
 457 -- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@).
 458 deleteAt :: Int -> Map k a -> Map k a
 459 deleteAt i map
 460   = updateAt (\k x -> Nothing) i map
 461
 462
 463 {--------------------------------------------------------------------
 464   Minimal, Maximal
 465 --------------------------------------------------------------------}
 466 -- | /O(log n)/. The minimal key of the map.
 467 findMin :: Map k a -> (k,a)
 468 findMin (Bin _ kx x Tip r)  = (kx,x)
 469 findMin (Bin _ kx x l r)    = findMin l
 470 findMin Tip                 = error "Map.findMin: empty tree has no minimal element"
 471
 472 -- | /O(log n)/. The maximal key of the map.
 473 findMax :: Map k a -> (k,a)
 474 findMax (Bin _ kx x l Tip)  = (kx,x)
 475 findMax (Bin _ kx x l r)    = findMax r
 476 findMax Tip                 = error "Map.findMax: empty tree has no maximal element"
 477
 478 -- | /O(log n)/. Delete the minimal key.
 479 deleteMin :: Map k a -> Map k a
 480 deleteMin (Bin _ kx x Tip r)  = r
 481 deleteMin (Bin _ kx x l r)    = balance kx x (deleteMin l) r
 482 deleteMin Tip                 = Tip
 483
 484 -- | /O(log n)/. Delete the maximal key.
 485 deleteMax :: Map k a -> Map k a
 486 deleteMax (Bin _ kx x l Tip)  = l
 487 deleteMax (Bin _ kx x l r)    = balance kx x l (deleteMax r)
 488 deleteMax Tip                 = Tip
 489
 490 -- | /O(log n)/. Update the value at the minimal key.
 491 updateMin :: (a -> Maybe a) -> Map k a -> Map k a
 492 updateMin f m
 493   = updateMinWithKey (\k x -> f x) m
 494
 495 -- | /O(log n)/. Update the value at the maximal key.
 496 updateMax :: (a -> Maybe a) -> Map k a -> Map k a
 497 updateMax f m
 498   = updateMaxWithKey (\k x -> f x) m
 499
 500
 501 -- | /O(log n)/. Update the value at the minimal key.
 502 updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
 503 updateMinWithKey f t
 504   = case t of
 505       Bin sx kx x Tip r  -> case f kx x of
 506                               Nothing -> r
 507                               Just x' -> Bin sx kx x' Tip r
 508       Bin sx kx x l r    -> balance kx x (updateMinWithKey f l) r
 509       Tip                -> Tip
 510
 511 -- | /O(log n)/. Update the value at the maximal key.
 512 updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
 513 updateMaxWithKey f t
 514   = case t of
 515       Bin sx kx x l Tip  -> case f kx x of
 516                               Nothing -> l
 517                               Just x' -> Bin sx kx x' l Tip
 518       Bin sx kx x l r    -> balance kx x l (updateMaxWithKey f r)
 519       Tip                -> Tip
 520
 521
 522 {--------------------------------------------------------------------
 523   Union.
 524 --------------------------------------------------------------------}
 525 -- | The union of a list of maps:
 526 --   (@'unions' == 'Prelude.foldl' 'union' 'empty'@).
 527 unions :: Ord k => [Map k a] -> Map k a
 528 unions ts
 529   = foldlStrict union empty ts
 530
 531 -- | The union of a list of maps, with a combining operation:
 532 --   (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@).
 533 unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a
 534 unionsWith f ts
 535   = foldlStrict (unionWith f) empty ts
 536
 537 -- | /O(n+m)/.
 538 -- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@.
 539 -- It prefers @t1@ when duplicate keys are encountered,
 540 -- i.e. (@'union' == 'unionWith' 'const'@).
 541 -- The implementation uses the efficient /hedge-union/ algorithm.
 542 -- Hedge-union is more efficient on (bigset `union` smallset)?
 543 union :: Ord k => Map k a -> Map k a -> Map k a
 544 union Tip t2  = t2
 545 union t1 Tip  = t1
 546 union t1 t2
 547    | size t1 >= size t2  = hedgeUnionL (const LT) (const GT) t1 t2
 548    | otherwise           = hedgeUnionR (const LT) (const GT) t2 t1
 549
 550 -- left-biased hedge union
 551 hedgeUnionL cmplo cmphi t1 Tip
 552   = t1
 553 hedgeUnionL cmplo cmphi Tip (Bin _ kx x l r)
 554   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 555 hedgeUnionL cmplo cmphi (Bin _ kx x l r) t2
 556   = join kx x (hedgeUnionL cmplo cmpkx l (trim cmplo cmpkx t2))
 557               (hedgeUnionL cmpkx cmphi r (trim cmpkx cmphi t2))
 558   where
 559     cmpkx k  = compare kx k
 560
 561 -- right-biased hedge union
 562 hedgeUnionR cmplo cmphi t1 Tip
 563   = t1
 564 hedgeUnionR cmplo cmphi Tip (Bin _ kx x l r)
 565   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 566 hedgeUnionR cmplo cmphi (Bin _ kx x l r) t2
 567   = join kx newx (hedgeUnionR cmplo cmpkx l lt)
 568                  (hedgeUnionR cmpkx cmphi r gt)
 569   where
 570     cmpkx k     = compare kx k
 571     lt          = trim cmplo cmpkx t2
 572     (found,gt)  = trimLookupLo kx cmphi t2
 573     newx        = case found of
 574                     Nothing -> x
 575                     Just y  -> y
 576
 577 {--------------------------------------------------------------------
 578   Union with a combining function
 579 --------------------------------------------------------------------}
 580 -- | /O(n+m)/. Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
 581 unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
 582 unionWith f m1 m2
 583   = unionWithKey (\k x y -> f x y) m1 m2
 584
 585 -- | /O(n+m)/.
 586 -- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
 587 -- Hedge-union is more efficient on (bigset `union` smallset).
 588 unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
 589 unionWithKey f Tip t2  = t2
 590 unionWithKey f t1 Tip  = t1
 591 unionWithKey f t1 t2
 592   | size t1 >= size t2  = hedgeUnionWithKey f (const LT) (const GT) t1 t2
 593   | otherwise           = hedgeUnionWithKey flipf (const LT) (const GT) t2 t1
 594   where
 595     flipf k x y   = f k y x
 596
 597 hedgeUnionWithKey f cmplo cmphi t1 Tip
 598   = t1
 599 hedgeUnionWithKey f cmplo cmphi Tip (Bin _ kx x l r)
 600   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 601 hedgeUnionWithKey f cmplo cmphi (Bin _ kx x l r) t2
 602   = join kx newx (hedgeUnionWithKey f cmplo cmpkx l lt)
 603                  (hedgeUnionWithKey f cmpkx cmphi r gt)
 604   where
 605     cmpkx k     = compare kx k
 606     lt          = trim cmplo cmpkx t2
 607     (found,gt)  = trimLookupLo kx cmphi t2
 608     newx        = case found of
 609                     Nothing -> x
 610                     Just y  -> f kx x y
 611
 612 {--------------------------------------------------------------------
 613   Difference
 614 --------------------------------------------------------------------}
 615 -- | /O(n+m)/. Difference of two maps.
 616 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 617 difference :: Ord k => Map k a -> Map k b -> Map k a
 618 difference Tip t2  = Tip
 619 difference t1 Tip  = t1
 620 difference t1 t2   = hedgeDiff (const LT) (const GT) t1 t2
 621
 622 hedgeDiff cmplo cmphi Tip t
 623   = Tip
 624 hedgeDiff cmplo cmphi (Bin _ kx x l r) Tip
 625   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 626 hedgeDiff cmplo cmphi t (Bin _ kx x l r)
 627   = merge (hedgeDiff cmplo cmpkx (trim cmplo cmpkx t) l)
 628           (hedgeDiff cmpkx cmphi (trim cmpkx cmphi t) r)
 629   where
 630     cmpkx k = compare kx k
 631
 632 -- | /O(n+m)/. Difference with a combining function.
 633 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 634 differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
 635 differenceWith f m1 m2
 636   = differenceWithKey (\k x y -> f x y) m1 m2
 637
 638 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
 639 -- encountered, the combining function is applied to the key and both values.
 640 -- If it returns 'Nothing', the element is discarded (proper set difference). If
 641 -- it returns (@'Just' y@), the element is updated with a new value @y@.
 642 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 643 differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
 644 differenceWithKey f Tip t2  = Tip
 645 differenceWithKey f t1 Tip  = t1
 646 differenceWithKey f t1 t2   = hedgeDiffWithKey f (const LT) (const GT) t1 t2
 647
 648 hedgeDiffWithKey f cmplo cmphi Tip t
 649   = Tip
 650 hedgeDiffWithKey f cmplo cmphi (Bin _ kx x l r) Tip
 651   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 652 hedgeDiffWithKey f cmplo cmphi t (Bin _ kx x l r)
 653   = case found of
 654       Nothing -> merge tl tr
 655       Just y  -> case f kx y x of
 656                    Nothing -> merge tl tr
 657                    Just z  -> join kx z tl tr
 658   where
 659     cmpkx k     = compare kx k
 660     lt          = trim cmplo cmpkx t
 661     (found,gt)  = trimLookupLo kx cmphi t
 662     tl          = hedgeDiffWithKey f cmplo cmpkx lt l
 663     tr          = hedgeDiffWithKey f cmpkx cmphi gt r
 664
 665
 666
 667 {--------------------------------------------------------------------
 668   Intersection
 669 --------------------------------------------------------------------}
 670 -- | /O(n+m)/. Intersection of two maps. The values in the first
 671 -- map are returned, i.e. (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@).
 672 intersection :: Ord k => Map k a -> Map k b -> Map k a
 673 intersection m1 m2
 674   = intersectionWithKey (\k x y -> x) m1 m2
 675
 676 -- | /O(n+m)/. Intersection with a combining function.
 677 intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c
 678 intersectionWith f m1 m2
 679   = intersectionWithKey (\k x y -> f x y) m1 m2
 680
 681 -- | /O(n+m)/. Intersection with a combining function.
 682 -- Intersection is more efficient on (bigset `intersection` smallset)
 683 intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
 684 intersectionWithKey f Tip t = Tip
 685 intersectionWithKey f t Tip = Tip
 686 intersectionWithKey f t1 t2
 687   | size t1 >= size t2  = intersectWithKey f t1 t2
 688   | otherwise           = intersectWithKey flipf t2 t1
 689   where
 690     flipf k x y   = f k y x
 691
 692 intersectWithKey f Tip t = Tip
 693 intersectWithKey f t Tip = Tip
 694 intersectWithKey f t (Bin _ kx x l r)
 695   = case found of
 696       Nothing -> merge tl tr
 697       Just y  -> join kx (f kx y x) tl tr
 698   where
 699     (lt,found,gt) = splitLookup kx t
 700     tl            = intersectWithKey f lt l
 701     tr            = intersectWithKey f gt r
 702
 703
 704
 705 {--------------------------------------------------------------------
 706   Submap
 707 --------------------------------------------------------------------}
 708 -- | /O(n+m)/.
 709 -- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
 710 isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
 711 isSubmapOf m1 m2
 712   = isSubmapOfBy (==) m1 m2
 713
 714 {- | /O(n+m)/.
 715  The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if
 716  all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when
 717  applied to their respective values. For example, the following
 718  expressions are all 'True':
 719
 720  > isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 721  > isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 722  > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)])
 723
 724  But the following are all 'False':
 725
 726  > isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)])
 727  > isSubmapOfBy (<)  (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 728  > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)])
 729 -}
 730 isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool
 731 isSubmapOfBy f t1 t2
 732   = (size t1 <= size t2) && (submap' f t1 t2)
 733
 734 submap' f Tip t = True
 735 submap' f t Tip = False
 736 submap' f (Bin _ kx x l r) t
 737   = case found of
 738       Nothing -> False
 739       Just y  -> f x y && submap' f l lt && submap' f r gt
 740   where
 741     (lt,found,gt) = splitLookup kx t
 742
 743 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
 744 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
 745 isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
 746 isProperSubmapOf m1 m2
 747   = isProperSubmapOfBy (==) m1 m2
 748
 749 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
 750  The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
 751  @m1@ and @m2@ are not equal,
 752  all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
 753  applied to their respective values. For example, the following
 754  expressions are all 'True':
 755
 756   > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 757   > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 758
 759  But the following are all 'False':
 760
 761   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
 762   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
 763   > isProperSubmapOfBy (<)  (fromList [(1,1)])       (fromList [(1,1),(2,2)])
 764 -}
 765 isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool
 766 isProperSubmapOfBy f t1 t2
 767   = (size t1 < size t2) && (submap' f t1 t2)
 768
 769 {--------------------------------------------------------------------
 770   Filter and partition
 771 --------------------------------------------------------------------}
 772 -- | /O(n)/. Filter all values that satisfy the predicate.
 773 filter :: Ord k => (a -> Bool) -> Map k a -> Map k a
 774 filter p m
 775   = filterWithKey (\k x -> p x) m
 776
 777 -- | /O(n)/. Filter all keys\/values that satisfy the predicate.
 778 filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a
 779 filterWithKey p Tip = Tip
 780 filterWithKey p (Bin _ kx x l r)
 781   | p kx x    = join kx x (filterWithKey p l) (filterWithKey p r)
 782   | otherwise = merge (filterWithKey p l) (filterWithKey p r)
 783
 784
 785 -- | /O(n)/. partition the map according to a predicate. The first
 786 -- map contains all elements that satisfy the predicate, the second all
 787 -- elements that fail the predicate. See also 'split'.
 788 partition :: Ord k => (a -> Bool) -> Map k a -> (Map k a,Map k a)
 789 partition p m
 790   = partitionWithKey (\k x -> p x) m
 791
 792 -- | /O(n)/. partition the map according to a predicate. The first
 793 -- map contains all elements that satisfy the predicate, the second all
 794 -- elements that fail the predicate. See also 'split'.
 795 partitionWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> (Map k a,Map k a)
 796 partitionWithKey p Tip = (Tip,Tip)
 797 partitionWithKey p (Bin _ kx x l r)
 798   | p kx x    = (join kx x l1 r1,merge l2 r2)
 799   | otherwise = (merge l1 r1,join kx x l2 r2)
 800   where
 801     (l1,l2) = partitionWithKey p l
 802     (r1,r2) = partitionWithKey p r
 803
 804
 805 {--------------------------------------------------------------------
 806   Mapping
 807 --------------------------------------------------------------------}
 808 -- | /O(n)/. Map a function over all values in the map.
 809 map :: (a -> b) -> Map k a -> Map k b
 810 map f m
 811   = mapWithKey (\k x -> f x) m
 812
 813 -- | /O(n)/. Map a function over all values in the map.
 814 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
 815 mapWithKey f Tip = Tip
 816 mapWithKey f (Bin sx kx x l r)
 817   = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
 818
 819 -- | /O(n)/. The function 'mapAccum' threads an accumulating
 820 -- argument through the map in ascending order of keys.
 821 mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 822 mapAccum f a m
 823   = mapAccumWithKey (\a k x -> f a x) a m
 824
 825 -- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating
 826 -- argument through the map in ascending order of keys.
 827 mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 828 mapAccumWithKey f a t
 829   = mapAccumL f a t
 830
 831 -- | /O(n)/. The function 'mapAccumL' threads an accumulating
 832 -- argument throught the map in ascending order of keys.
 833 mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 834 mapAccumL f a t
 835   = case t of
 836       Tip -> (a,Tip)
 837       Bin sx kx x l r
 838           -> let (a1,l') = mapAccumL f a l
 839                  (a2,x') = f a1 kx x
 840                  (a3,r') = mapAccumL f a2 r
 841              in (a3,Bin sx kx x' l' r')
 842
 843 -- | /O(n)/. The function 'mapAccumR' threads an accumulating
 844 -- argument throught the map in descending order of keys.
 845 mapAccumR :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 846 mapAccumR f a t
 847   = case t of
 848       Tip -> (a,Tip)
 849       Bin sx kx x l r
 850           -> let (a1,r') = mapAccumR f a r
 851                  (a2,x') = f a1 kx x
 852                  (a3,l') = mapAccumR f a2 l
 853              in (a3,Bin sx kx x' l' r')
 854
 855 -- | /O(n*log n)/.
 856 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
 857 --
 858 -- The size of the result may be smaller if @f@ maps two or more distinct
 859 -- keys to the same new key.  In this case the value at the smallest of
 860 -- these keys is retained.
 861
 862 mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a
 863 mapKeys = mapKeysWith (\x y->x)
 864
 865 -- | /O(n*log n)/.
 866 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
 867 --
 868 -- The size of the result may be smaller if @f@ maps two or more distinct
 869 -- keys to the same new key.  In this case the associated values will be
 870 -- combined using @c@.
 871
 872 mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a
 873 mapKeysWith c f = fromListWith c . List.map fFirst . toList
 874     where fFirst (x,y) = (f x, y)
 875
 876
 877 -- | /O(n)/.
 878 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
 879 -- is strictly monotonic.
 880 -- /The precondition is not checked./
 881 -- Semi-formally, we have:
 882 --
 883 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
 884 -- >                     ==> mapKeysMonotonic f s == mapKeys f s
 885 -- >     where ls = keys s
 886
 887 mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a
 888 mapKeysMonotonic f Tip = Tip
 889 mapKeysMonotonic f (Bin sz k x l r) =
 890     Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r)
 891
 892 {--------------------------------------------------------------------
 893   Folds
 894 --------------------------------------------------------------------}
 895
 896 -- | /O(n)/. Fold the values in the map, such that
 897 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
 898 -- For example,
 899 --
 900 -- > elems map = fold (:) [] map
 901 --
 902 fold :: (a -> b -> b) -> b -> Map k a -> b
 903 fold f z m
 904   = foldWithKey (\k x z -> f x z) z m
 905
 906 -- | /O(n)/. Fold the keys and values in the map, such that
 907 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
 908 -- For example,
 909 --
 910 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
 911 --
 912 foldWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
 913 foldWithKey f z t
 914   = foldr f z t
 915
 916 -- | /O(n)/. In-order fold.
 917 foldi :: (k -> a -> b -> b -> b) -> b -> Map k a -> b
 918 foldi f z Tip               = z
 919 foldi f z (Bin _ kx x l r)  = f kx x (foldi f z l) (foldi f z r)
 920
 921 -- | /O(n)/. Post-order fold.
 922 foldr :: (k -> a -> b -> b) -> b -> Map k a -> b
 923 foldr f z Tip              = z
 924 foldr f z (Bin _ kx x l r) = foldr f (f kx x (foldr f z r)) l
 925
 926 -- | /O(n)/. Pre-order fold.
 927 foldl :: (b -> k -> a -> b) -> b -> Map k a -> b
 928 foldl f z Tip              = z
 929 foldl f z (Bin _ kx x l r) = foldl f (f (foldl f z l) kx x) r
 930
 931 {--------------------------------------------------------------------
 932   List variations
 933 --------------------------------------------------------------------}
 934 -- | /O(n)/.
 935 -- Return all elements of the map in the ascending order of their keys.
 936 elems :: Map k a -> [a]
 937 elems m
 938   = [x | (k,x) <- assocs m]
 939
 940 -- | /O(n)/. Return all keys of the map in ascending order.
 941 keys  :: Map k a -> [k]
 942 keys m
 943   = [k | (k,x) <- assocs m]
 944
 945 -- | /O(n)/. The set of all keys of the map.
 946 keysSet :: Map k a -> Set.Set k
 947 keysSet m = Set.fromDistinctAscList (keys m)
 948
 949 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
 950 assocs :: Map k a -> [(k,a)]
 951 assocs m
 952   = toList m
 953
 954 {--------------------------------------------------------------------
 955   Lists
 956   use [foldlStrict] to reduce demand on the control-stack
 957 --------------------------------------------------------------------}
 958 -- | /O(n*log n)/. Build a map from a list of key\/value pairs. See also 'fromAscList'.
 959 fromList :: Ord k => [(k,a)] -> Map k a
 960 fromList xs
 961   = foldlStrict ins empty xs
 962   where
 963     ins t (k,x) = insert k x t
 964
 965 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
 966 fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a
 967 fromListWith f xs
 968   = fromListWithKey (\k x y -> f x y) xs
 969
 970 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
 971 fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
 972 fromListWithKey f xs
 973   = foldlStrict ins empty xs
 974   where
 975     ins t (k,x) = insertWithKey f k x t
 976
 977 -- | /O(n)/. Convert to a list of key\/value pairs.
 978 toList :: Map k a -> [(k,a)]
 979 toList t      = toAscList t
 980
 981 -- | /O(n)/. Convert to an ascending list.
 982 toAscList :: Map k a -> [(k,a)]
 983 toAscList t   = foldr (\k x xs -> (k,x):xs) [] t
 984
 985 -- | /O(n)/.
 986 toDescList :: Map k a -> [(k,a)]
 987 toDescList t  = foldl (\xs k x -> (k,x):xs) [] t
 988
 989
 990 {--------------------------------------------------------------------
 991   Building trees from ascending/descending lists can be done in linear time.
 992
 993   Note that if [xs] is ascending that:
 994     fromAscList xs       == fromList xs
 995     fromAscListWith f xs == fromListWith f xs
 996 --------------------------------------------------------------------}
 997 -- | /O(n)/. Build a map from an ascending list in linear time.
 998 -- /The precondition (input list is ascending) is not checked./
 999 fromAscList :: Eq k => [(k,a)] -> Map k a
1000 fromAscList xs
1001   = fromAscListWithKey (\k x y -> x) xs
1002
1003 -- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys.
1004 -- /The precondition (input list is ascending) is not checked./
1005 fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
1006 fromAscListWith f xs
1007   = fromAscListWithKey (\k x y -> f x y) xs
1008
1009 -- | /O(n)/. Build a map from an ascending list in linear time with a
1010 -- combining function for equal keys.
1011 -- /The precondition (input list is ascending) is not checked./
1012 fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
1013 fromAscListWithKey f xs
1014   = fromDistinctAscList (combineEq f xs)
1015   where
1016   -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1017   combineEq f xs
1018     = case xs of
1019         []     -> []
1020         [x]    -> [x]
1021         (x:xx) -> combineEq' x xx
1022
1023   combineEq' z [] = [z]
1024   combineEq' z@(kz,zz) (x@(kx,xx):xs)
1025     | kx==kz    = let yy = f kx xx zz in combineEq' (kx,yy) xs
1026     | otherwise = z:combineEq' x xs
1027
1028
1029 -- | /O(n)/. Build a map from an ascending list of distinct elements in linear time.
1030 -- /The precondition is not checked./
1031 fromDistinctAscList :: [(k,a)] -> Map k a
1032 fromDistinctAscList xs
1033   = build const (length xs) xs
1034   where
1035     -- 1) use continutations so that we use heap space instead of stack space.
1036     -- 2) special case for n==5 to build bushier trees.
1037     build c 0 xs   = c Tip xs
1038     build c 5 xs   = case xs of
1039                        ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx)
1040                             -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx
1041     build c n xs   = seq nr $ build (buildR nr c) nl xs
1042                    where
1043                      nl = n `div` 2
1044                      nr = n - nl - 1
1045
1046     buildR n c l ((k,x):ys) = build (buildB l k x c) n ys
1047     buildB l k x c r zs     = c (bin k x l r) zs
1048
1049
1050
1051 {--------------------------------------------------------------------
1052   Utility functions that return sub-ranges of the original
1053   tree. Some functions take a comparison function as argument to
1054   allow comparisons against infinite values. A function [cmplo k]
1055   should be read as [compare lo k].
1056
1057   [trim cmplo cmphi t]  A tree that is either empty or where [cmplo k == LT]
1058                         and [cmphi k == GT] for the key [k] of the root.
1059   [filterGt cmp t]      A tree where for all keys [k]. [cmp k == LT]
1060   [filterLt cmp t]      A tree where for all keys [k]. [cmp k == GT]
1061
1062   [split k t]           Returns two trees [l] and [r] where all keys
1063                         in [l] are <[k] and all keys in [r] are >[k].
1064   [splitLookup k t]     Just like [split] but also returns whether [k]
1065                         was found in the tree.
1066 --------------------------------------------------------------------}
1067
1068 {--------------------------------------------------------------------
1069   [trim lo hi t] trims away all subtrees that surely contain no
1070   values between the range [lo] to [hi]. The returned tree is either
1071   empty or the key of the root is between @lo@ and @hi@.
1072 --------------------------------------------------------------------}
1073 trim :: (k -> Ordering) -> (k -> Ordering) -> Map k a -> Map k a
1074 trim cmplo cmphi Tip = Tip
1075 trim cmplo cmphi t@(Bin sx kx x l r)
1076   = case cmplo kx of
1077       LT -> case cmphi kx of
1078               GT -> t
1079               le -> trim cmplo cmphi l
1080       ge -> trim cmplo cmphi r
1081
1082 trimLookupLo :: Ord k => k -> (k -> Ordering) -> Map k a -> (Maybe a, Map k a)
1083 trimLookupLo lo cmphi Tip = (Nothing,Tip)
1084 trimLookupLo lo cmphi t@(Bin sx kx x l r)
1085   = case compare lo kx of
1086       LT -> case cmphi kx of
1087               GT -> (lookup lo t, t)
1088               le -> trimLookupLo lo cmphi l
1089       GT -> trimLookupLo lo cmphi r
1090       EQ -> (Just x,trim (compare lo) cmphi r)
1091
1092
1093 {--------------------------------------------------------------------
1094   [filterGt k t] filter all keys >[k] from tree [t]
1095   [filterLt k t] filter all keys <[k] from tree [t]
1096 --------------------------------------------------------------------}
1097 filterGt :: Ord k => (k -> Ordering) -> Map k a -> Map k a
1098 filterGt cmp Tip = Tip
1099 filterGt cmp (Bin sx kx x l r)
1100   = case cmp kx of
1101       LT -> join kx x (filterGt cmp l) r
1102       GT -> filterGt cmp r
1103       EQ -> r
1104
1105 filterLt :: Ord k => (k -> Ordering) -> Map k a -> Map k a
1106 filterLt cmp Tip = Tip
1107 filterLt cmp (Bin sx kx x l r)
1108   = case cmp kx of
1109       LT -> filterLt cmp l
1110       GT -> join kx x l (filterLt cmp r)
1111       EQ -> l
1112
1113 {--------------------------------------------------------------------
1114   Split
1115 --------------------------------------------------------------------}
1116 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@ where
1117 -- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
1118 split :: Ord k => k -> Map k a -> (Map k a,Map k a)
1119 split k Tip = (Tip,Tip)
1120 split k (Bin sx kx x l r)
1121   = case compare k kx of
1122       LT -> let (lt,gt) = split k l in (lt,join kx x gt r)
1123       GT -> let (lt,gt) = split k r in (join kx x l lt,gt)
1124       EQ -> (l,r)
1125
1126 -- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just
1127 -- like 'split' but also returns @'lookup' k map@.
1128 splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a)
1129 splitLookup k Tip = (Tip,Nothing,Tip)
1130 splitLookup k (Bin sx kx x l r)
1131   = case compare k kx of
1132       LT -> let (lt,z,gt) = splitLookup k l in (lt,z,join kx x gt r)
1133       GT -> let (lt,z,gt) = splitLookup k r in (join kx x l lt,z,gt)
1134       EQ -> (l,Just x,r)
1135
1136 {--------------------------------------------------------------------
1137   Utility functions that maintain the balance properties of the tree.
1138   All constructors assume that all values in [l] < [k] and all values
1139   in [r] > [k], and that [l] and [r] are valid trees.
1140
1141   In order of sophistication:
1142     [Bin sz k x l r]  The type constructor.
1143     [bin k x l r]     Maintains the correct size, assumes that both [l]
1144                       and [r] are balanced with respect to each other.
1145     [balance k x l r] Restores the balance and size.
1146                       Assumes that the original tree was balanced and
1147                       that [l] or [r] has changed by at most one element.
1148     [join k x l r]    Restores balance and size.
1149
1150   Furthermore, we can construct a new tree from two trees. Both operations
1151   assume that all values in [l] < all values in [r] and that [l] and [r]
1152   are valid:
1153     [glue l r]        Glues [l] and [r] together. Assumes that [l] and
1154                       [r] are already balanced with respect to each other.
1155     [merge l r]       Merges two trees and restores balance.
1156
1157   Note: in contrast to Adam's paper, we use (<=) comparisons instead
1158   of (<) comparisons in [join], [merge] and [balance].
1159   Quickcheck (on [difference]) showed that this was necessary in order
1160   to maintain the invariants. It is quite unsatisfactory that I haven't
1161   been able to find out why this is actually the case! Fortunately, it
1162   doesn't hurt to be a bit more conservative.
1163 --------------------------------------------------------------------}
1164
1165 {--------------------------------------------------------------------
1166   Join
1167 --------------------------------------------------------------------}
1168 join :: Ord k => k -> a -> Map k a -> Map k a -> Map k a
1169 join kx x Tip r  = insertMin kx x r
1170 join kx x l Tip  = insertMax kx x l
1171 join kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz)
1172   | delta*sizeL <= sizeR  = balance kz z (join kx x l lz) rz
1173   | delta*sizeR <= sizeL  = balance ky y ly (join kx x ry r)
1174   | otherwise             = bin kx x l r
1175
1176
1177 -- insertMin and insertMax don't perform potentially expensive comparisons.
1178 insertMax,insertMin :: k -> a -> Map k a -> Map k a
1179 insertMax kx x t
1180   = case t of
1181       Tip -> singleton kx x
1182       Bin sz ky y l r
1183           -> balance ky y l (insertMax kx x r)
1184
1185 insertMin kx x t
1186   = case t of
1187       Tip -> singleton kx x
1188       Bin sz ky y l r
1189           -> balance ky y (insertMin kx x l) r
1190
1191 {--------------------------------------------------------------------
1192   [merge l r]: merges two trees.
1193 --------------------------------------------------------------------}
1194 merge :: Map k a -> Map k a -> Map k a
1195 merge Tip r   = r
1196 merge l Tip   = l
1197 merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry)
1198   | delta*sizeL <= sizeR = balance ky y (merge l ly) ry
1199   | delta*sizeR <= sizeL = balance kx x lx (merge rx r)
1200   | otherwise            = glue l r
1201
1202 {--------------------------------------------------------------------
1203   [glue l r]: glues two trees together.
1204   Assumes that [l] and [r] are already balanced with respect to each other.
1205 --------------------------------------------------------------------}
1206 glue :: Map k a -> Map k a -> Map k a
1207 glue Tip r = r
1208 glue l Tip = l
1209 glue l r
1210   | size l > size r = let ((km,m),l') = deleteFindMax l in balance km m l' r
1211   | otherwise       = let ((km,m),r') = deleteFindMin r in balance km m l r'
1212
1213
1214 -- | /O(log n)/. Delete and find the minimal element.
1215 deleteFindMin :: Map k a -> ((k,a),Map k a)
1216 deleteFindMin t
1217   = case t of
1218       Bin _ k x Tip r -> ((k,x),r)
1219       Bin _ k x l r   -> let (km,l') = deleteFindMin l in (km,balance k x l' r)
1220       Tip             -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip)
1221
1222 -- | /O(log n)/. Delete and find the maximal element.
1223 deleteFindMax :: Map k a -> ((k,a),Map k a)
1224 deleteFindMax t
1225   = case t of
1226       Bin _ k x l Tip -> ((k,x),l)
1227       Bin _ k x l r   -> let (km,r') = deleteFindMax r in (km,balance k x l r')
1228       Tip             -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip)
1229
1230
1231 {--------------------------------------------------------------------
1232   [balance l x r] balances two trees with value x.
1233   The sizes of the trees should balance after decreasing the
1234   size of one of them. (a rotation).
1235
1236   [delta] is the maximal relative difference between the sizes of
1237           two trees, it corresponds with the [w] in Adams' paper.
1238   [ratio] is the ratio between an outer and inner sibling of the
1239           heavier subtree in an unbalanced setting. It determines
1240           whether a double or single rotation should be performed
1241           to restore balance. It is correspondes with the inverse
1242           of $\alpha$ in Adam's article.
1243
1244   Note that:
1245   - [delta] should be larger than 4.646 with a [ratio] of 2.
1246   - [delta] should be larger than 3.745 with a [ratio] of 1.534.
1247
1248   - A lower [delta] leads to a more 'perfectly' balanced tree.
1249   - A higher [delta] performs less rebalancing.
1250
1251   - Balancing is automatic for random data and a balancing
1252     scheme is only necessary to avoid pathological worst cases.
1253     Almost any choice will do, and in practice, a rather large
1254     [delta] may perform better than smaller one.
1255
1256   Note: in contrast to Adam's paper, we use a ratio of (at least) [2]
1257   to decide whether a single or double rotation is needed. Allthough
1258   he actually proves that this ratio is needed to maintain the
1259   invariants, his implementation uses an invalid ratio of [1].
1260 --------------------------------------------------------------------}
1261 delta,ratio :: Int
1262 delta = 5
1263 ratio = 2
1264
1265 balance :: k -> a -> Map k a -> Map k a -> Map k a
1266 balance k x l r
1267   | sizeL + sizeR <= 1    = Bin sizeX k x l r
1268   | sizeR >= delta*sizeL  = rotateL k x l r
1269   | sizeL >= delta*sizeR  = rotateR k x l r
1270   | otherwise             = Bin sizeX k x l r
1271   where
1272     sizeL = size l
1273     sizeR = size r
1274     sizeX = sizeL + sizeR + 1
1275
1276 -- rotate
1277 rotateL k x l r@(Bin _ _ _ ly ry)
1278   | size ly < ratio*size ry = singleL k x l r
1279   | otherwise               = doubleL k x l r
1280
1281 rotateR k x l@(Bin _ _ _ ly ry) r
1282   | size ry < ratio*size ly = singleR k x l r
1283   | otherwise               = doubleR k x l r
1284
1285 -- basic rotations
1286 singleL k1 x1 t1 (Bin _ k2 x2 t2 t3)  = bin k2 x2 (bin k1 x1 t1 t2) t3
1287 singleR k1 x1 (Bin _ k2 x2 t1 t2) t3  = bin k2 x2 t1 (bin k1 x1 t2 t3)
1288
1289 doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4)
1290 doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4)
1291
1292
1293 {--------------------------------------------------------------------
1294   The bin constructor maintains the size of the tree
1295 --------------------------------------------------------------------}
1296 bin :: k -> a -> Map k a -> Map k a -> Map k a
1297 bin k x l r
1298   = Bin (size l + size r + 1) k x l r
1299
1300
1301 {--------------------------------------------------------------------
1302   Eq converts the tree to a list. In a lazy setting, this
1303   actually seems one of the faster methods to compare two trees
1304   and it is certainly the simplest :-)
1305 --------------------------------------------------------------------}
1306 instance (Eq k,Eq a) => Eq (Map k a) where
1307   t1 == t2  = (size t1 == size t2) && (toAscList t1 == toAscList t2)
1308
1309 {--------------------------------------------------------------------
1310   Ord
1311 --------------------------------------------------------------------}
1312
1313 instance (Ord k, Ord v) => Ord (Map k v) where
1314     compare m1 m2 = compare (toAscList m1) (toAscList m2)
1315
1316 {--------------------------------------------------------------------
1317   Functor
1318 --------------------------------------------------------------------}
1319 instance Functor (Map k) where
1320   fmap f m  = map f m
1321
1322 {--------------------------------------------------------------------
1323   Read
1324 --------------------------------------------------------------------}
1325 instance (Ord k, Read k, Read e) => Read (Map k e) where
1326 #ifdef __GLASGOW_HASKELL__
1327   readPrec = parens $ prec 10 $ do
1328     Ident "fromList" <- lexP
1329     xs <- readPrec
1330     return (fromList xs)
1331
1332   readListPrec = readListPrecDefault
1333 #else
1334   readsPrec p = readParen (p > 10) $ \ r -> do
1335     ("fromList",s) <- lex r
1336     (xs,t) <- reads s
1337     return (fromList xs,t)
1338 #endif
1339
1340 -- parses a pair of things with the syntax a:=b
1341 readPair :: (Read a, Read b) => ReadS (a,b)
1342 readPair s = do (a, ct1)    <- reads s
1343                 (":=", ct2) <- lex ct1
1344                 (b, ct3)    <- reads ct2
1345                 return ((a,b), ct3)
1346
1347 {--------------------------------------------------------------------
1348   Show
1349 --------------------------------------------------------------------}
1350 instance (Show k, Show a) => Show (Map k a) where
1351   showsPrec d m  = showParen (d > 10) $
1352     showString "fromList " . shows (toList m)
1353
1354 showMap :: (Show k,Show a) => [(k,a)] -> ShowS
1355 showMap []
1356   = showString "{}"
1357 showMap (x:xs)
1358   = showChar '{' . showElem x . showTail xs
1359   where
1360     showTail []     = showChar '}'
1361     showTail (x:xs) = showString ", " . showElem x . showTail xs
1362
1363     showElem (k,x)  = shows k . showString " := " . shows x
1364
1365
1366 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1367 -- in a compressed, hanging format.
1368 showTree :: (Show k,Show a) => Map k a -> String
1369 showTree m
1370   = showTreeWith showElem True False m
1371   where
1372     showElem k x  = show k ++ ":=" ++ show x
1373
1374
1375 {- | /O(n)/. The expression (@'showTreeWith' showelem hang wide map@) shows
1376  the tree that implements the map. Elements are shown using the @showElem@ function. If @hang@ is
1377  'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1378  @wide@ is 'True', an extra wide version is shown.
1379
1380 >  Map> let t = fromDistinctAscList [(x,()) | x <- [1..5]]
1381 >  Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True False t
1382 >  (4,())
1383 >  +--(2,())
1384 >  |  +--(1,())
1385 >  |  +--(3,())
1386 >  +--(5,())
1387 >
1388 >  Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True True t
1389 >  (4,())
1390 >  |
1391 >  +--(2,())
1392 >  |  |
1393 >  |  +--(1,())
1394 >  |  |
1395 >  |  +--(3,())
1396 >  |
1397 >  +--(5,())
1398 >
1399 >  Map> putStrLn $ showTreeWith (\k x -> show (k,x)) False True t
1400 >  +--(5,())
1401 >  |
1402 >  (4,())
1403 >  |
1404 >  |  +--(3,())
1405 >  |  |
1406 >  +--(2,())
1407 >     |
1408 >     +--(1,())
1409
1410 -}
1411 showTreeWith :: (k -> a -> String) -> Bool -> Bool -> Map k a -> String
1412 showTreeWith showelem hang wide t
1413   | hang      = (showsTreeHang showelem wide [] t) ""
1414   | otherwise = (showsTree showelem wide [] [] t) ""
1415
1416 showsTree :: (k -> a -> String) -> Bool -> [String] -> [String] -> Map k a -> ShowS
1417 showsTree showelem wide lbars rbars t
1418   = case t of
1419       Tip -> showsBars lbars . showString "|\n"
1420       Bin sz kx x Tip Tip
1421           -> showsBars lbars . showString (showelem kx x) . showString "\n"
1422       Bin sz kx x l r
1423           -> showsTree showelem wide (withBar rbars) (withEmpty rbars) r .
1424              showWide wide rbars .
1425              showsBars lbars . showString (showelem kx x) . showString "\n" .
1426              showWide wide lbars .
1427              showsTree showelem wide (withEmpty lbars) (withBar lbars) l
1428
1429 showsTreeHang :: (k -> a -> String) -> Bool -> [String] -> Map k a -> ShowS
1430 showsTreeHang showelem wide bars t
1431   = case t of
1432       Tip -> showsBars bars . showString "|\n"
1433       Bin sz kx x Tip Tip
1434           -> showsBars bars . showString (showelem kx x) . showString "\n"
1435       Bin sz kx x l r
1436           -> showsBars bars . showString (showelem kx x) . showString "\n" .
1437              showWide wide bars .
1438              showsTreeHang showelem wide (withBar bars) l .
1439              showWide wide bars .
1440              showsTreeHang showelem wide (withEmpty bars) r
1441
1442
1443 showWide wide bars
1444   | wide      = showString (concat (reverse bars)) . showString "|\n"
1445   | otherwise = id
1446
1447 showsBars :: [String] -> ShowS
1448 showsBars bars
1449   = case bars of
1450       [] -> id
1451       _  -> showString (concat (reverse (tail bars))) . showString node
1452
1453 node           = "+--"
1454 withBar bars   = "|  ":bars
1455 withEmpty bars = "   ":bars
1456
1457 {--------------------------------------------------------------------
1458   Typeable
1459 --------------------------------------------------------------------}
1460
1461 #include "Typeable.h"
1462 INSTANCE_TYPEABLE2(Map,mapTc,"Map")
1463
1464 {--------------------------------------------------------------------
1465   Assertions
1466 --------------------------------------------------------------------}
1467 -- | /O(n)/. Test if the internal map structure is valid.
1468 valid :: Ord k => Map k a -> Bool
1469 valid t
1470   = balanced t && ordered t && validsize t
1471
1472 ordered t
1473   = bounded (const True) (const True) t
1474   where
1475     bounded lo hi t
1476       = case t of
1477           Tip              -> True
1478           Bin sz kx x l r  -> (lo kx) && (hi kx) && bounded lo (<kx) l && bounded (>kx) hi r
1479
1480 -- | Exported only for "Debug.QuickCheck"
1481 balanced :: Map k a -> Bool
1482 balanced t
1483   = case t of
1484       Tip              -> True
1485       Bin sz kx x l r  -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) &&
1486                           balanced l && balanced r
1487
1488
1489 validsize t
1490   = (realsize t == Just (size t))
1491   where
1492     realsize t
1493       = case t of
1494           Tip             -> Just 0
1495           Bin sz kx x l r -> case (realsize l,realsize r) of
1496                               (Just n,Just m)  | n+m+1 == sz  -> Just sz
1497                               other            -> Nothing
1498
1499 {--------------------------------------------------------------------
1500   Utilities
1501 --------------------------------------------------------------------}
1502 foldlStrict f z xs
1503   = case xs of
1504       []     -> z
1505       (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1506
1507
1508 {-
1509 {--------------------------------------------------------------------
1510   Testing
1511 --------------------------------------------------------------------}
1512 testTree xs   = fromList [(x,"*") | x <- xs]
1513 test1 = testTree [1..20]
1514 test2 = testTree [30,29..10]
1515 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1516
1517 {--------------------------------------------------------------------
1518   QuickCheck
1519 --------------------------------------------------------------------}
1520 qcheck prop
1521   = check config prop
1522   where
1523     config = Config
1524       { configMaxTest = 500
1525       , configMaxFail = 5000
1526       , configSize    = \n -> (div n 2 + 3)
1527       , configEvery   = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1528       }
1529
1530
1531 {--------------------------------------------------------------------
1532   Arbitrary, reasonably balanced trees
1533 --------------------------------------------------------------------}
1534 instance (Enum k,Arbitrary a) => Arbitrary (Map k a) where
1535   arbitrary = sized (arbtree 0 maxkey)
1536             where maxkey  = 10000
1537
1538 arbtree :: (Enum k,Arbitrary a) => Int -> Int -> Int -> Gen (Map k a)
1539 arbtree lo hi n
1540   | n <= 0        = return Tip
1541   | lo >= hi      = return Tip
1542   | otherwise     = do{ x  <- arbitrary
1543                       ; i  <- choose (lo,hi)
1544                       ; m  <- choose (1,30)
1545                       ; let (ml,mr)  | m==(1::Int)= (1,2)
1546                                      | m==2       = (2,1)
1547                                      | m==3       = (1,1)
1548                                      | otherwise  = (2,2)
1549                       ; l  <- arbtree lo (i-1) (n `div` ml)
1550                       ; r  <- arbtree (i+1) hi (n `div` mr)
1551                       ; return (bin (toEnum i) x l r)
1552                       }
1553
1554
1555 {--------------------------------------------------------------------
1556   Valid tree's
1557 --------------------------------------------------------------------}
1558 forValid :: (Show k,Enum k,Show a,Arbitrary a,Testable b) => (Map k a -> b) -> Property
1559 forValid f
1560   = forAll arbitrary $ \t ->
1561 --    classify (balanced t) "balanced" $
1562     classify (size t == 0) "empty" $
1563     classify (size t > 0  && size t <= 10) "small" $
1564     classify (size t > 10 && size t <= 64) "medium" $
1565     classify (size t > 64) "large" $
1566     balanced t ==> f t
1567
1568 forValidIntTree :: Testable a => (Map Int Int -> a) -> Property
1569 forValidIntTree f
1570   = forValid f
1571
1572 forValidUnitTree :: Testable a => (Map Int () -> a) -> Property
1573 forValidUnitTree f
1574   = forValid f
1575
1576
1577 prop_Valid
1578   = forValidUnitTree $ \t -> valid t
1579
1580 {--------------------------------------------------------------------
1581   Single, Insert, Delete
1582 --------------------------------------------------------------------}
1583 prop_Single :: Int -> Int -> Bool
1584 prop_Single k x
1585   = (insert k x empty == singleton k x)
1586
1587 prop_InsertValid :: Int -> Property
1588 prop_InsertValid k
1589   = forValidUnitTree $ \t -> valid (insert k () t)
1590
1591 prop_InsertDelete :: Int -> Map Int () -> Property
1592 prop_InsertDelete k t
1593   = (lookup k t == Nothing) ==> delete k (insert k () t) == t
1594
1595 prop_DeleteValid :: Int -> Property
1596 prop_DeleteValid k
1597   = forValidUnitTree $ \t ->
1598     valid (delete k (insert k () t))
1599
1600 {--------------------------------------------------------------------
1601   Balance
1602 --------------------------------------------------------------------}
1603 prop_Join :: Int -> Property
1604 prop_Join k
1605   = forValidUnitTree $ \t ->
1606     let (l,r) = split k t
1607     in valid (join k () l r)
1608
1609 prop_Merge :: Int -> Property
1610 prop_Merge k
1611   = forValidUnitTree $ \t ->
1612     let (l,r) = split k t
1613     in valid (merge l r)
1614
1615
1616 {--------------------------------------------------------------------
1617   Union
1618 --------------------------------------------------------------------}
1619 prop_UnionValid :: Property
1620 prop_UnionValid
1621   = forValidUnitTree $ \t1 ->
1622     forValidUnitTree $ \t2 ->
1623     valid (union t1 t2)
1624
1625 prop_UnionInsert :: Int -> Int -> Map Int Int -> Bool
1626 prop_UnionInsert k x t
1627   = union (singleton k x) t == insert k x t
1628
1629 prop_UnionAssoc :: Map Int Int -> Map Int Int -> Map Int Int -> Bool
1630 prop_UnionAssoc t1 t2 t3
1631   = union t1 (union t2 t3) == union (union t1 t2) t3
1632
1633 prop_UnionComm :: Map Int Int -> Map Int Int -> Bool
1634 prop_UnionComm t1 t2
1635   = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1636
1637 prop_UnionWithValid
1638   = forValidIntTree $ \t1 ->
1639     forValidIntTree $ \t2 ->
1640     valid (unionWithKey (\k x y -> x+y) t1 t2)
1641
1642 prop_UnionWith :: [(Int,Int)] -> [(Int,Int)] -> Bool
1643 prop_UnionWith xs ys
1644   = sum (elems (unionWith (+) (fromListWith (+) xs) (fromListWith (+) ys)))
1645     == (sum (Prelude.map snd xs) + sum (Prelude.map snd ys))
1646
1647 prop_DiffValid
1648   = forValidUnitTree $ \t1 ->
1649     forValidUnitTree $ \t2 ->
1650     valid (difference t1 t2)
1651
1652 prop_Diff :: [(Int,Int)] -> [(Int,Int)] -> Bool
1653 prop_Diff xs ys
1654   =  List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1655     == List.sort ((List.\\) (nub (Prelude.map fst xs))  (nub (Prelude.map fst ys)))
1656
1657 prop_IntValid
1658   = forValidUnitTree $ \t1 ->
1659     forValidUnitTree $ \t2 ->
1660     valid (intersection t1 t2)
1661
1662 prop_Int :: [(Int,Int)] -> [(Int,Int)] -> Bool
1663 prop_Int xs ys
1664   =  List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1665     == List.sort (nub ((List.intersect) (Prelude.map fst xs)  (Prelude.map fst ys)))
1666
1667 {--------------------------------------------------------------------
1668   Lists
1669 --------------------------------------------------------------------}
1670 prop_Ordered
1671   = forAll (choose (5,100)) $ \n ->
1672     let xs = [(x,()) | x <- [0..n::Int]]
1673     in fromAscList xs == fromList xs
1674
1675 prop_List :: [Int] -> Bool
1676 prop_List xs
1677   = (sort (nub xs) == [x | (x,()) <- toList (fromList [(x,()) | x <- xs])])
1678 -}