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