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