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