Data/Map.hs

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