Data/Map.hs

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